Update app.py

app.py

@@ -1,3 +1,74 @@
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
+from threading import Thread
+from queue import Empty
+
+# Load the tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("thrishala/mental_health_chatbot")
+model = AutoModelForCausalLM.from_pretrained("thrishala/mental_health_chatbot", torch_dtype=torch.float16)
+
+# Move model to GPU if available
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = model.to(device)
+
+class StopOnTokens(StoppingCriteria):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        stop_ids = [29, 0]  # Token IDs for stopping criteria
+        for stop_id in stop_ids:
+            if input_ids[0][-1] == stop_id:
+                return True
+        return False
+
+def predict(message, history):
+    # Prepare the input history in the expected format for the model
+    history_transformer_format = list(zip(history[:-1], history[1:])) + [[message, ""]]
+    stop = StopOnTokens()
+
+    # Concatenate conversation history
+    messages = "".join(["".join(["\n<human>:"+item[0], "\n<bot>:"+item[1]]) for item in history_transformer_format])
+
+    # Tokenize and prepare model inputs
+    model_inputs = tokenizer([messages], return_tensors="pt").to(device)
+
+    # Create streamer with longer timeout
+    streamer = TextIteratorStreamer(tokenizer, timeout=30., skip_prompt=True, skip_special_tokens=True)
+    
+    # Define generation parameters
+    generate_kwargs = dict(
+        model_inputs,
+        streamer=streamer,
+        max_new_tokens=512,  # Reduced to avoid memory issues
+        do_sample=True,
+        top_p=0.85,  # Adjusted for faster generation
+        top_k=500,  # Adjusted for faster generation
+        temperature=1.0,
+        num_beams=1,
+        stopping_criteria=StoppingCriteriaList([stop])
+    )
+
+    # Run the generation in a separate thread
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+
+    # Yield generated tokens
+    partial_message = ""
+    try:
+        for new_token in streamer:
+            print(f"Received token: {new_token}")  # Debugging output
+            if new_token != '<':
+                partial_message += new_token
+                yield partial_message
+    except Empty:
+        print("No tokens were generated within the timeout period.")
+
+# Gradio interface to run the chatbot
+gr.ChatInterface(predict).launch()
+
+
+
+
+
 
 
 # import gradio as gr
@@ -162,66 +233,5 @@
 
 
 
-import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
-from threading import Thread
-
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("thrishala/mental_health_chatbot")
-
-# Check if CUDA (GPU) is available, otherwise use CPU
-device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
-model = AutoModelForCausalLM.from_pretrained("thrishala/mental_health_chatbot", torch_dtype=torch.float16)
-model = model.to(device)
-
-# Custom stopping criteria to stop generation on specific tokens
-class StopOnTokens(StoppingCriteria):
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-        stop_ids = [29, 0]  # EOS token or any other token you want to stop on
-        for stop_id in stop_ids:
-            if input_ids[0][-1] == stop_id:
-                return True
-        return False
-
-def predict(message, history):
-    # Prepare the message history for the model
-    history_transformer_format = list(zip(history[:-1], history[1:])) + [[message, ""]]
-    stop = StopOnTokens()
-
-    # Format the conversation for the model
-    messages = "".join(["".join(["\n<human>:"+item[0], "\n<bot>:"+item[1]]) for item in history_transformer_format])
-
-    # Tokenize input and move to the correct device (GPU or CPU)
-    model_inputs = tokenizer([messages], return_tensors="pt").to(device)
-    
-    # Create a streamer to handle model outputs
-    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
-
-    # Generation parameters
-    generate_kwargs = dict(
-        model_inputs,
-        streamer=streamer,
-        max_new_tokens=1024,
-        do_sample=True,
-        top_p=0.95,
-        top_k=1000,
-        temperature=1.0,
-        num_beams=1,
-        stopping_criteria=StoppingCriteriaList([stop])
-    )
-    
-    # Run the generation in a separate thread
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    # Collect the generated tokens
-    partial_message = ""
-    for new_token in streamer:
-        if new_token != '<':  # Avoid issues with special tokens
-            partial_message += new_token
-            yield partial_message
 
-# Launch the Gradio interface
-gr.ChatInterface(predict).launch()