Update app.py

app.py

@@ -1,23 +1,23 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import spaces
-import threading
 
-model_name = "INSAIT-Institute/MamayLM-Gemma-2-9B-IT-v0.1"
+
+model_name = "INSAIT-Institute/MamayLM-Gemma-2-9B-IT-v0.1"  
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
 
 @spaces.GPU
 def respond(message, chat_history, system_message, max_new_tokens, temperature, top_p):
+
     prompt = f"{system_message.strip()}\n"
     for user, bot in chat_history:
         prompt += f"User: {user}\nAssistant: {bot}\n"
     prompt += f"User: {message}\nAssistant:"
 
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = dict(
+    output = model.generate(
         **inputs,
         max_new_tokens=int(max_new_tokens),
         pad_token_id=tokenizer.eos_token_id,
@@ -25,17 +25,11 @@ def respond(message, chat_history, system_message, max_new_tokens, temperature,
         temperature=float(temperature),
         top_p=float(top_p),
         eos_token_id=tokenizer.eos_token_id,
-        streamer=streamer,
     )
+    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
 
-    thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-
-    partial_text = ""
-    for new_text in streamer:
-        partial_text += new_text
-        response = partial_text.split("User:")[0].strip()
-        yield response
+    response = decoded.split("Assistant:")[-1].strip().split("User:")[0].strip()
+    return response
 
 def clear_fn():
     return None
@@ -60,6 +54,7 @@ chat = gr.ChatInterface(
     ],
     title="💬 Chat with MamayLM",
     description="A multi-turn chat interface for MamayLM-v0.1-9B with configurable parameters.",
-    theme="soft",
+    theme="soft"
 )
+
 chat.launch()