Update app.py

app.py

@@ -1,39 +1,23 @@
-import os
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-import re
-import gradio as gr
-import spaces
-
-model_id = "tatihden/gemma_mental_health_2b_it_en"
-dtype = torch.bfloat16
+from transformers import pipeline
+from transformers import Conversation
 
-@spaces.GPU
-def gemma_chat(message, history):
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        hidden_activation="gelu_pytorch_tanh",
-        device_map="cuda",
-        torch_dtype=dtype,
-    )
+import gradio as gr
 
-    chat = [
-        { "role": "user", "content": message },
-    ]
+chatbot = pipeline(model="hf://tatihden/gemma_mental_health_2b_it_en")
 
-    prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
-    
-    inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
-    outputs = model.generate(input_ids=inputs.to(model.device), max_new_tokens=2048)
-    
-    response = tokenizer.decode(outputs[0])
+message_list = []
+response_list = []
 
-    # clean the response
-    response_cleaned = re.split("model", response)
+def mini_chatbot(message, history):
+    conversation = Conversation(text=message,
+                                past_user_inputs=message_list,
+                                generated_responses=response_list)
+    conversation = chatbot(conversation)
 
-    # return the response
-    return response_cleaned
+    return conversation.generated_responses[-1]
 
+demo_chatbot = gr.ChatInterface(mini_chatbot,
+                                title="CalmChat",
+                                description="Enter text to start chatting.")
 
-gr.ChatInterface(gemma_chat).launch()
+demo_chatbot.launch()