Update app.py

app.py

@@ -25,19 +25,32 @@ def respond(message, history):
     # Add current message
     prompt += f"User: {message}\nAssistant: "
     
-    # Generate response
-    response = pipe(
-        prompt,
-        max_new_tokens=10000,
-        temperature=0.7,
-        do_sample=True,
-        pad_token_id=pipe.tokenizer.eos_token_id,
-        return_full_text=False,
-    )
+    # Generate response with streaming
+    streamer = pipe.tokenizer.decode
     
-    # Extract and yield the generated text
-    generated_text = response[0]["generated_text"]
-    yield generated_text
+    # Generate tokens
+    inputs = pipe.tokenizer(prompt, return_tensors="pt").to(pipe.model.device)
+    
+    with torch.no_grad():
+        outputs = pipe.model.generate(
+            **inputs,
+            max_new_tokens=10000,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=pipe.tokenizer.eos_token_id,
+            return_full_text=False,
+        )
+    
+    # Decode the generated tokens, skipping the input tokens
+    generated_tokens = outputs[0][inputs['input_ids'].shape[-1]:]
+    
+    # Stream the output token by token
+    response_text = ""
+    for i in range(len(generated_tokens)):
+        token = generated_tokens[i:i+1]
+        token_text = pipe.tokenizer.decode(token, skip_special_tokens=True)
+        response_text += token_text
+        yield response_text
 
 # Create the chat interface
 demo = gr.ChatInterface(