Update app.py

app.py

@@ -35,6 +35,7 @@ class ChatModel:
                 config["model_name"],
                 device_map="auto",
                 torch_dtype=torch.float16,
+                attn_implementation="flash_attention_2" if "phi-3" in model_name else "eager",
                 trust_remote_code=True
             )
 
@@ -47,9 +48,9 @@ class ChatModel:
         prompt = config["template"].format(message=message)
         inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
         
-        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
+        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, timeout=60)
         generation_kwargs = dict(
-            inputs.input_ids,
+            **inputs,
             streamer=streamer,
             max_new_tokens=512,
             temperature=0.7,
@@ -67,24 +68,24 @@ class ChatModel:
 model_handler = ChatModel()
 
 def chat(message, history, model_choice):
-    # Initialize streaming
-    streamer, tokenizer, start_time = model_handler.stream_response(message, model_choice)
-    buffer = ""
-    
-    # Stream tokens
-    for new_text in streamer:
-        buffer += new_text
-        yield [(message, buffer)]
-    
-    # Add performance metrics
-    elapsed_time = time.time() - start_time
-    tokens = len(tokenizer.encode(buffer))
-    token_speed = tokens / elapsed_time if elapsed_time > 0 else 0
-    final_response = f"{buffer}\n\n⏱️ {elapsed_time:.2f}s | 🚀 {token_speed:.2f} tokens/s"
-    yield [(message, final_response)]
+    try:
+        streamer, tokenizer, start_time = model_handler.stream_response(message, model_choice)
+        buffer = ""
+        
+        for new_text in streamer:
+            buffer += new_text
+            yield [(message, buffer)]
+        
+        elapsed_time = time.time() - start_time
+        tokens = len(tokenizer.encode(buffer))
+        token_speed = tokens / elapsed_time if elapsed_time > 0 else 0
+        final_response = f"{buffer}\n\n⏱️ {elapsed_time:.2f}s | 🚀 {token_speed:.2f} tokens/s"
+        yield [(message, final_response)]
+    except Exception as e:
+        yield [(message, f"Error: {str(e)}")]
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚀 Streaming LLM Chatbot")
+    gr.Markdown("# 🚀 Streaming LLM Chatbot (Fixed)")
     with gr.Row():
         model_choice = gr.Dropdown(
             choices=["phi-3", "llama3-8b"],