r1 i want it better

app.py

@@ -2,27 +2,32 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-if torch.cuda.is_available():
-    device = "cuda"
-else:
-    device = "cpu"
+# Determine device
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 model_id = "thrishala/mental_health_chatbot"
 
 try:
+    # Load model with appropriate device_map and settings
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
-        device_map=device,
+        device_map="auto",  # Use "auto" for device_map instead of device name
         torch_dtype=torch.float16,
         low_cpu_mem_usage=True,
-        max_memory={device: "15GB"},
+        max_memory={0: "15GiB"} if torch.cuda.is_available() else None,
         offload_folder="offload",
-    )
+    ).eval()  # Set model to evaluation mode
+
     tokenizer = AutoTokenizer.from_pretrained(model_id)
-    tokenizer.model_max_length = 512
+    tokenizer.pad_token = tokenizer.eos_token  # Set padding token if missing
 
+    # Perform a dummy generation to initialize model (if needed)
     dummy_input = tokenizer("This is a test.", return_tensors="pt").to(model.device)
-    model.generate(input_ids=dummy_input.input_ids, return_dict=True)  # Dummy call
+    model.generate(
+        input_ids=dummy_input.input_ids,
+        max_new_tokens=1,
+        return_dict_in_generate=True  # Correct parameter name
+    )
 
 except Exception as e:
     print(f"Error loading model: {e}")
@@ -37,32 +42,34 @@ def generate_text(prompt, max_new_tokens=128):
             max_new_tokens=max_new_tokens,
             do_sample=False,
             eos_token_id=tokenizer.eos_token_id,
-            return_dict=True,  # Explicitly set return_dict=True
+            return_dict_in_generate=True  # Correct parameter name
         )
 
-    generated_text = tokenizer.decode(output.sequences[0], skip_special_tokens=True)  # Decode from sequences
+    generated_text = tokenizer.decode(output.sequences[0], skip_special_tokens=True)
     return generated_text
 
 def generate_text_streaming(prompt, max_new_tokens=128):
     input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
 
     with torch.no_grad():
-        for i in range(max_new_tokens):
+        for _ in range(max_new_tokens):
             output = model.generate(
                 input_ids=input_ids,
                 max_new_tokens=1,
                 do_sample=False,
                 eos_token_id=tokenizer.eos_token_id,
-                return_dict=True,
-                output_scores=True,
+                return_dict_in_generate=True  # Correct parameter name
             )
 
-            generated_token = tokenizer.decode(output.logits[0][-1].argmax(), skip_special_tokens=True)
+            # Get the last generated token
+            generated_token_id = output.sequences[0, -1]
+            generated_token = tokenizer.decode([generated_token_id], skip_special_tokens=True)
             yield generated_token
 
+            # Append new token to input_ids
             input_ids = torch.cat([input_ids, output.sequences[:, -1:]], dim=-1)
 
-            if output.sequences[0][-1] == tokenizer.eos_token_id:
+            if generated_token_id == tokenizer.eos_token_id:
                 break
 
 def respond(message, history, system_message, max_tokens):
@@ -73,7 +80,7 @@ def respond(message, history, system_message, max_tokens):
 
     try:
         for token in generate_text_streaming(prompt, max_tokens):
-            yield token  # Yield each token individually
+            yield token
     except Exception as e:
         print(f"Error during generation: {e}")
         yield "An error occurred."
@@ -90,4 +97,4 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()