Update app.py

app.py

@@ -2,71 +2,51 @@ 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)
 
+tokenizer = AutoTokenizer.from_pretrained("togethercomputer/RedPajama-INCITE-Chat-3B-v1")
+model = AutoModelForCausalLM.from_pretrained("togethercomputer/RedPajama-INCITE-Chat-3B-v1", 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
+        stop_ids = [29, 0]
         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)
+    messages = "".join(["".join(["\n<human>:"+item[0], "\n<bot>:"+item[1]])
+                for item in history_transformer_format])
 
-    # Create streamer with longer timeout
-    streamer = TextIteratorStreamer(tokenizer, timeout=30., skip_prompt=True, skip_special_tokens=True)
-    
-    # Define generation parameters
+    model_inputs = tokenizer([messages], return_tensors="pt").to("cuda")
+    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
         model_inputs,
         streamer=streamer,
-        max_new_tokens=512,  # Reduced to avoid memory issues
+        max_new_tokens=1024,
         do_sample=True,
-        top_p=0.85,  # Adjusted for faster generation
-        top_k=500,  # Adjusted for faster generation
+        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()
 
-    # Yield generated tokens
     partial_message = ""
-    try:
-        for new_token in streamer:
-            if new_token != '<':
-                partial_message += new_token
-                yield partial_message
-    except Empty:
-        yield "Error: No tokens generated or generation timeout."
-
-# Gradio interface to run the chatbot
-gr.ChatInterface(predict).launch()
-
-
-
+    for new_token in streamer:
+        if new_token != '<':
+            partial_message += new_token
+            yield partial_message
 
+gr.ChatInterface(predict).launch()