more

app.py

@@ -19,14 +19,16 @@ FastLanguageModel.for_inference(model)  # Enable optimized inference
 
 # Define the response function
 def respond(message, history, system_message, max_tokens, temperature, top_p):
-    # Combine system and user inputs
-    messages = [{"role": "system", "content": system_message}] + [
-        {"role": "user", "content": user_msg} if assistant_msg is None else {"role": "assistant", "content": assistant_msg}
-        for user_msg, assistant_msg in history
-    ]
+    # Combine system message and conversation history
+    messages = [{"role": "system", "content": system_message}]
+    for user_msg, assistant_msg in history:
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if assistant_msg:
+            messages.append({"role": "assistant", "content": assistant_msg})
     messages.append({"role": "user", "content": message})
     
-    # Apply the chat template
+    # Tokenize inputs
     inputs = tokenizer.apply_chat_template(
         messages,
         tokenize=True,
@@ -34,8 +36,21 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
         return_tensors="pt",
     ).to("cuda" if torch.cuda.is_available() else "cpu")
 
-    # Use a TextStreamer for real-time decoding
-    streamer = TextStreamer(tokenizer, skip_prompt=True)
+    # Use TextStreamer to process and yield outputs incrementally
+    class GradioStreamer(TextStreamer):
+        def __init__(self, tokenizer, *args, **kwargs):
+            super().__init__(tokenizer, *args, **kwargs)
+            self.generated_text = ""
+
+        def on_token(self, token_id):
+            token = self.tokenizer.decode(token_id, skip_special_tokens=True)
+            self.generated_text += token
+            yield self.generated_text
+
+    # Initialize Gradio-compatible streamer
+    streamer = GradioStreamer(tokenizer, skip_prompt=True)
+    
+    # Generate response with streaming
     _ = model.generate(
         input_ids=inputs,
         max_new_tokens=max_tokens,