Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
@@ -7,29 +7,30 @@ print("Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 print("Loading model on CPU...")
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto")
-model.to("cpu")
+model = AutoModelForCausalLM.from_pretrained(model_id)
+model = model.to("cpu")
 
 pipe = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    max_new_tokens=256,
+    device=-1,  # ensures CPU mode
+    max_new_tokens=200,
     temperature=0.7,
     top_p=0.9,
-    repetition_penalty=1.05
 )
 
 def chat_fn(message, history):
     prompt = f"User: {message}\nAssistant:"
     output = pipe(prompt, num_return_sequences=1)
-    reply = output[0]['generated_text'].split("Assistant:")[-1].strip()
-    return reply
+    text = output[0]["generated_text"]
+    response = text.split("Assistant:")[-1].strip()
+    return response
 
 with gr.Blocks() as demo:
     gr.Markdown("<h1 style='text-align:center;'>🧠 Mental Health Companion Bot</h1>")
     chatbot = gr.Chatbot()
-    msg = gr.Textbox(placeholder="Type your message...")
+    msg = gr.Textbox(placeholder="Type your message here...")
     clear = gr.Button("Clear Chat")
 
     def respond(message, chat_history):