Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
 import torch
+import time  # Added for timing
 
 MODEL_CONFIG = {
     "phi-3": {
@@ -48,6 +49,7 @@ class ChatModel:
             self.tokenizers[model_name] = tokenizer
 
     def generate(self, message, model_name, history):
+        start_time = time.time()  # Start timing
         self.load_model(model_name)
         config = MODEL_CONFIG[model_name]
         
@@ -67,20 +69,27 @@ class ChatModel:
             return_full_text=False
         )
         
-        response = pipe(prompt)[0]['generated_text']
-        return response.strip()
+        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
+        
+        return response, elapsed_time, tokens_per_sec
 
 model_handler = ChatModel()
 
 def chat(message, history, model_choice):
     try:
-        response = model_handler.generate(message, model_choice, history)
-        return [(message, response)]
+        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)}")]
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚀 Phi-3 vs Llama-3 Chatbot")
+    gr.Markdown("# 🚀 LLM Chatbot with Performance Metrics")
     with gr.Row():
         model_choice = gr.Dropdown(
             choices=["phi-3", "llama3-8b"],