Update app.py

app.py

@@ -1,7 +1,6 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftConfig, PeftModel
-from threading import Thread
 import gradio as gr
 import spaces
 
@@ -25,7 +24,6 @@ def load_model():
                 BASE_MODEL,
                 torch_dtype=torch.float16,
                 device_map="auto",
-                load_in_4bit=True,
                 trust_remote_code=True
             )
             
@@ -43,29 +41,27 @@ def load_model():
 def generate_response(prompt, max_new_tokens=128):
     model, tokenizer = load_model()
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = dict(
-        input_ids=inputs.input_ids,
+    outputs = model.generate(
+        **inputs,
         max_new_tokens=max_new_tokens,
         temperature=0.7,
         top_p=0.9,
         repetition_penalty=1.2,
-        streamer=streamer,
+        do_sample=True
     )
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-    return streamer
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 def chat_with_zephyr(message, history):
     conversation_history = history[-3:]  # Limit to last 3 exchanges
     full_prompt = "\n".join([f"Human: {h[0]}\nZephyr: {h[1]}" for h in conversation_history])
     full_prompt += f"\nHuman: {message}\nZephyr:"
     
-    streamer = generate_response(full_prompt)
-    response = ""
-    for new_text in streamer:
-        response += new_text
-        yield response
+    response = generate_response(full_prompt)
+    
+    # Extract Zephyr's response
+    zephyr_response = response.split("Zephyr:")[-1].strip()
+    
+    return zephyr_response
 
 css = """
 body {