Update Gradio app with multiple files

app.py

@@ -1,42 +1,85 @@
 import gradio as gr
-from models import chat_with_model
-from config import MODEL_NAME
+import spaces
+from models import load_model, chat_with_model
+from config import AVAILABLE_MODELS, DEFAULT_MODEL
+from utils import format_conversation, save_chat_history, load_chat_history
 
-def chat_response(message, history):
-    # Format history for the model
-    conversation = []
-    for user_msg, assistant_msg in history:
-        conversation.append({"role": "user", "content": user_msg})
-        if assistant_msg:
-            conversation.append({"role": "assistant", "content": assistant_msg})
-    conversation.append({"role": "user", "content": message})
-    
-    response = chat_with_model(conversation)
+# Global variable for current model
+current_model = None
+
+def switch_model(model_name):
+    global current_model
+    current_model = load_model(model_name)
+    return f"Switched to {model_name}"
+
+@spaces.GPU(duration=120)  # GPU acceleration for inference
+def generate_response(message, history, model_name, streaming=False):
+    global current_model
+    if current_model is None or current_model['name'] != model_name:
+        current_model = load_model(model_name)
+    
+    conversation = format_conversation(message, history)
+    response = chat_with_model(current_model, conversation, streaming=streaming)
     return response
 
-with gr.Blocks(title="LLM Chatbot") as demo:
-    gr.Markdown("# 🤖 LLM Chatbot")
-    gr.Markdown("*Powered by a Hugging Face model*")
+with gr.Blocks(title="Advanced LLM Chatbot", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🤖 Advanced LLM Chatbot")
+    gr.Markdown("*Powered by Hugging Face models with model switching and streaming*")
     gr.Markdown("[Built with anycoder](https://huggingface.co/spaces/akhaliq/anycoder)")
     
-    chatbot = gr.Chatbot(type="messages", height=400)
+    with gr.Row():
+        model_dropdown = gr.Dropdown(
+            choices=list(AVAILABLE_MODELS.keys()), 
+            value=DEFAULT_MODEL, 
+            label="Select Model",
+            info="Choose a chat model to use"
+        )
+        switch_btn = gr.Button("Switch Model")
+        status = gr.Textbox(label="Status", interactive=False, value=f"Current: {DEFAULT_MODEL}")
+    
+    with gr.Row():
+        streaming_toggle = gr.Checkbox(label="Enable Streaming", value=False)
+        save_btn = gr.Button("Save Chat")
+        load_btn = gr.Button("Load Chat")
+    
+    chatbot = gr.Chatbot(type="messages", height=400, show_copy_button=True)
     msg = gr.Textbox(placeholder="Type your message here...", label="Message")
     
     with gr.Row():
-        submit_btn = gr.Button("Send")
+        submit_btn = gr.Button("Send", variant="primary")
         clear_btn = gr.ClearButton([msg, chatbot])
+        undo_btn = gr.Button("Undo Last")
     
-    def respond(message, chat_history):
+    def respond(message, chat_history, model_name, streaming):
         if not message.strip():
             return "", chat_history
         
-        response = chat_response(message, chat_history)
+        response = generate_response(message, chat_history, model_name, streaming)
         chat_history.append({"role": "user", "content": message})
         chat_history.append({"role": "assistant", "content": response})
         return "", chat_history
     
-    msg.submit(respond, [msg, chatbot], [msg, chatbot])
-    submit_btn.click(respond, [msg, chatbot], [msg, chatbot])
+    def undo_last(history):
+        if len(history) >= 2:
+            return history[:-2]  # Remove last user and assistant message
+        return history
+    
+    def save_history(history):
+        save_chat_history(history)
+        return "Chat saved!"
+    
+    def load_history():
+        return load_chat_history()
+    
+    model_dropdown.change(switch_model, model_dropdown, status)
+    switch_btn.click(switch_model, model_dropdown, status)
+    
+    msg.submit(respond, [msg, chatbot, model_dropdown, streaming_toggle], [msg, chatbot])
+    submit_btn.click(respond, [msg, chatbot, model_dropdown, streaming_toggle], [msg, chatbot])
+    
+    undo_btn.click(undo_last, chatbot, chatbot)
+    save_btn.click(save_history, chatbot, status)
+    load_btn.click(load_history, None, chatbot)
 
 if __name__ == "__main__":
     demo.launch()

config.py

@@ -1,2 +1,9 @@
 # Configuration for the chatbot
-MODEL_NAME = "microsoft/DialoGPT-medium"  # Change to your preferred Hugging Face chat model
+AVAILABLE_MODELS = {
+    "microsoft/DialoGPT-medium": "DialoGPT Medium (Conversational)",
+    "microsoft/DialoGPT-large": "DialoGPT Large (Advanced)",
+    "facebook/blenderbot-400M-distill": "BlenderBot 400M (Conversational)",
+    "microsoft/GODEL-v1_1-base-seq2seq": "GODEL Base (Goal-oriented)"
+}
+
+DEFAULT_MODEL = "microsoft/DialoGPT-medium"

models.py

@@ -1,22 +1,48 @@
-from transformers import pipeline
-from utils import format_conversation
-from config import MODEL_NAME
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# Load the model pipeline once
-chat_pipeline = None
+# Cache for loaded models
+model_cache = {}
 
-def get_chat_pipeline():
-    global chat_pipeline
-    if chat_pipeline is None:
-        chat_pipeline = pipeline("conversational", model=MODEL_NAME)
-    return chat_pipeline
+def load_model(model_name):
+    """Load and cache a Hugging Face model."""
+    if model_name not in model_cache:
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name, 
+                torch_dtype=torch.float16,  # Use float16 for efficiency
+                device_map="auto"  # Auto-detect GPU
+            )
+            model_cache[model_name] = {
+                'name': model_name,
+                'tokenizer': tokenizer,
+                'model': model
+            }
+        except Exception as e:
+            raise ValueError(f"Failed to load model {model_name}: {str(e)}")
+    
+    return model_cache[model_name]
 
-def chat_with_model(conversation):
-    """Send conversation to the Hugging Face model and get response."""
+def chat_with_model(model_data, conversation, streaming=False):
+    """Generate response using the loaded model."""
     try:
-        pipe = get_chat_pipeline()
-        formatted_input = format_conversation(conversation)
-        response = pipe(formatted_input)
-        return response[0]['generated_text']
+        tokenizer = model_data['tokenizer']
+        model = model_data['model']
+        
+        inputs = tokenizer(conversation, return_tensors="pt").to(model.device)
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_length=inputs['input_ids'].shape[1] + 100,  # Generate up to 100 new tokens
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.9,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+        return response.strip()
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"Error generating response: {str(e)}"

requirements.txt

@@ -1,4 +1,6 @@
 gradio>=4.0.0
-transformers>=4.0.0
+transformers>=4.21.0
 torch>=2.0.0
-accelerate>=0.20.0
+accelerate>=0.20.0
+spaces>=0.25.0  # For GPU acceleration on Spaces
+sentencepiece>=0.1.99  # Required for some tokenizers

utils.py

@@ -1,3 +1,22 @@
-def format_conversation(conversation):
+import json
+import os
+
+def format_conversation(message, history):
     """Format conversation for the model."""
-    return "\n".join([f"{msg['role']}: {msg['content']}" for msg in conversation])
+    conversation = []
+    for msg in history:
+        conversation.append(f"{msg['role']}: {msg['content']}")
+    conversation.append(f"user: {message}")
+    return "\n".join(conversation)
+
+def save_chat_history(history):
+    """Save chat history to a file."""
+    with open("chat_history.json", "w") as f:
+        json.dump(history, f)
+
+def load_chat_history():
+    """Load chat history from a file."""
+    if os.path.exists("chat_history.json"):
+        with open("chat_history.json", "r") as f:
+            return json.load(f)
+    return []