Update app.py

app.py

@@ -5,69 +5,55 @@ import torch
 model_id = "thrishala/mental_health_chatbot"
 
 try:
-    # Load model with int8 quantization for CPU
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         device_map="cpu",
-        torch_dtype=torch.float16,  # Use float16 for reduced memory
-        low_cpu_mem_usage=True,     # Enable memory optimization
-        max_memory={"cpu": "15GB"},  # Limit memory usage
-        offload_folder="offload",    # Enable disk offloading
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True,
+        max_memory={"cpu": "15GB"},
+        offload_folder="offload",
     )
     
-    # Load tokenizer
     tokenizer = AutoTokenizer.from_pretrained(model_id)
+    tokenizer.model_max_length = 256  # Set maximum length
     
-    # Create pipeline with optimizations
     pipe = pipeline(
         "text-generation",
         model=model,
         tokenizer=tokenizer,
         torch_dtype=torch.float16,
-        num_return_sequences=1,  # Only generate one response
-        do_sample=True,         # Enable sampling since we're using temperature and top_p
-        truncation=True,        # Explicitly enable truncation
-        max_new_tokens=128      # Use only max_new_tokens
+        num_return_sequences=1,
+        do_sample=True,
+        truncation=True,
+        max_new_tokens=128
     )
 
 except Exception as e:
     print(f"Error loading model: {e}")
     exit()
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,  # You can use this for initial instructions
-    max_tokens,
-    temperature,
-    top_p,
-):
-    # 2. Construct the Prompt (Crucial!)
-    prompt = f"{system_message}\n" 
+def respond(message, history, system_message, max_tokens, temperature, top_p):
+    prompt = f"{system_message}\n"
     for user_msg, bot_msg in history:
         prompt += f"User: {user_msg}\nAssistant: {bot_msg}\n"
     prompt += f"User: {message}\nAssistant:"
-    
-    # 3. Generate with the Pipeline
+
     try:
         response = pipe(
             prompt,
             max_new_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
         )[0]["generated_text"]
-        prompt,
-        do_sample=True,
-        #Extract the bot's reply (adjust if your model format is different)
+        
         bot_response = response.split("Assistant:")[-1].strip()
         yield bot_response
-    
     except Exception as e:
         print(f"Error during generation: {e}")
-        yield "An error occurred during generation." #Handle generation errors.
-
+        yield "An error occurred during generation."
 
-# 4. Gradio Interface (No changes needed here)
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
@@ -78,13 +64,10 @@ demo = gr.ChatInterface(
         gr.Slider(minimum=1, maximum=128, value=128, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
+            minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)",
         ),
     ],
+    chatbot=gr.Chatbot(type="messages"),  # Updated to new format
 )
 
 if __name__ == "__main__":