Update app.py

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from threading import Thread
 import torch
-import time  # Added for timing
+import time
 
 MODEL_CONFIG = {
     "phi-3": {
@@ -18,13 +19,6 @@ MODEL_CONFIG = {
     }
 }
 
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_use_double_quant=True
-)
-
 class ChatModel:
     def __init__(self):
         self.models = {}
@@ -34,62 +28,63 @@ class ChatModel:
         if model_name not in self.models:
             config = MODEL_CONFIG[model_name]
             
-            tokenizer = AutoTokenizer.from_pretrained(config["model_name"])
-            tokenizer.pad_token = tokenizer.eos_token
+            self.tokenizers[model_name] = AutoTokenizer.from_pretrained(config["model_name"])
+            self.tokenizers[model_name].pad_token = self.tokenizers[model_name].eos_token
             
-            model = AutoModelForCausalLM.from_pretrained(
+            self.models[model_name] = AutoModelForCausalLM.from_pretrained(
                 config["model_name"],
-                quantization_config=bnb_config,
                 device_map="auto",
                 torch_dtype=torch.float16,
                 trust_remote_code=True
             )
-            
-            self.models[model_name] = model
-            self.tokenizers[model_name] = tokenizer
 
-    def generate(self, message, model_name, history):
-        start_time = time.time()  # Start timing
+    def stream_response(self, message, model_name):
         self.load_model(model_name)
         config = MODEL_CONFIG[model_name]
+        tokenizer = self.tokenizers[model_name]
+        model = self.models[model_name]
         
-        # Format prompt
         prompt = config["template"].format(message=message)
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
         
-        # Create pipeline
-        pipe = pipeline(
-            "text-generation",
-            model=self.models[model_name],
-            tokenizer=self.tokenizers[model_name],
-            max_new_tokens=384,
+        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
+        generation_kwargs = dict(
+            inputs.input_ids,
+            streamer=streamer,
+            max_new_tokens=512,
             temperature=0.7,
             top_p=0.9,
             repetition_penalty=1.1,
             do_sample=True,
-            return_full_text=False
+            pad_token_id=tokenizer.eos_token_id
         )
         
-        response = pipe(prompt)[0]['generated_text'].strip()
-        
-        # Calculate metrics
-        elapsed_time = time.time() - start_time
-        tokens = len(self.tokenizers[model_name].encode(response))
-        tokens_per_sec = tokens / elapsed_time if elapsed_time > 0 else 0
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
         
-        return response, elapsed_time, tokens_per_sec
+        return streamer, tokenizer, time.time()
 
 model_handler = ChatModel()
 
 def chat(message, history, model_choice):
-    try:
-        response, response_time, token_speed = model_handler.generate(message, model_choice, history)
-        formatted_response = f"{response}\n\n⏱️ Response Time: {response_time:.2f}s | 🚀 Speed: {token_speed:.2f} tokens/s"
-        return [(message, formatted_response)]
-    except Exception as e:
-        return [(message, f"Error: {str(e)}")]
+    # 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)]
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚀 LLM Chatbot with Performance Metrics")
+    gr.Markdown("# 🚀 Streaming LLM Chatbot")
     with gr.Row():
         model_choice = gr.Dropdown(
             choices=["phi-3", "llama3-8b"],