test

app.py

@@ -7,55 +7,153 @@ import spaces
 loaded_models = {}
 
 # List of available models (update with your preferred models)
-models = ["hanzla/gemma-2b-datascience-instruct-v5", "hanzla/gemma-2b-datascience-instruct-v4.5"]
+models = [
+    "hanzla/gemma-2b-datascience-instruct-v5",
+    "hanzla/gemma-2b-datascience-instruct-v4.5"
+]
+
 
 @spaces.GPU
-def load_model(model_name):
-    if model_name not in loaded_models:
-        print(f"Loading model: {model_name}")
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda" if torch.cuda.is_available() else "cpu")
-        loaded_models[model_name] = (model, tokenizer)
-    return loaded_models[model_name]
+def load_all_models():
+    """
+    Pre-loads all models and their tokenizers into memory.
+    """
+    for model_name in models:
+        if model_name not in loaded_models:
+            print(f"Loading model: {model_name}")
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda" if torch.cuda.is_available() else "cpu")
+            loaded_models[model_name] = (model, tokenizer)
+    return "All models loaded successfully."
+
 
 @spaces.GPU
 def get_model_response(model_name, message):
-    model, tokenizer = load_model(model_name)
+    """
+    Generates a response from the specified model given a user message.
+    """
+    model, tokenizer = loaded_models[model_name]
     inputs = tokenizer(message, return_tensors="pt").to(model.device)
 
-    outputs = model.generate(**inputs)
+    # Generate response with appropriate parameters
+    outputs = model.generate(
+        **inputs,
+        max_length=512,
+        do_sample=True,
+        top_p=0.95,
+        top_k=50
+    )
 
-    response = tokenizer.decode(outputs[0])
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return response
 
 
-def chat(message, history, model1, model2):
+def chat(message, history1, history2, model1, model2):
+    """
+    Handles the chat interaction by getting responses from both models
+    and updating their respective histories.
+    """
     response1 = get_model_response(model1, message)
     response2 = get_model_response(model2, message)
-    history = history or []
-    history.append({"role": "user", "content": message})
-    history.append({"role": "assistant", "content": f"{model1}: {response1}\n\n{model2}: {response2}"})
-    return history
+
+    history1 = history1 or []
+    history2 = history2 or []
+
+    # Update history for Model 1
+    history1.append(("User", message))
+    history1.append((model1.split("/")[-1], response1))
+
+    # Update history for Model 2
+    history2.append(("User", message))
+    history2.append((model2.split("/")[-1], response2))
+
+    return history1, history2
+
+
+# Initialize vote counts
+vote_counts = {"model1": 0, "model2": 0}
 
 
+def upvote_vote(model1, model2):
+    """
+    Increments the vote count for Model 1 and returns updated counts.
+    """
+    vote_counts["model1"] += 1
+    return f"Votes - {model1.split('/')[-1]}: {vote_counts['model1']}, {model2.split('/')[-1]}: {vote_counts['model2']}"
+
+
+def downvote_vote(model1, model2):
+    """
+    Increments the vote count for Model 2 and returns updated counts.
+    """
+    vote_counts["model2"] += 1
+    return f"Votes - {model1.split('/')[-1]}: {vote_counts['model1']}, {model2.split('/')[-1]}: {vote_counts['model2']}"
+
+
+def clear_chat():
+    """
+    Clears both chat histories and resets vote counts.
+    """
+    global vote_counts
+    vote_counts = {"model1": 0, "model2": 0}
+    return [], [], "Votes - 0, 0"
+
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Hugging Face Model Comparison Chat")
+    gr.Markdown("# 🤖 Hugging Face Model Comparison Chat")
 
+    # Dropdowns for selecting models
     with gr.Row():
         model1_dropdown = gr.Dropdown(choices=models, label="Model 1", value=models[0])
         model2_dropdown = gr.Dropdown(choices=models, label="Model 2", value=models[1])
 
-    chatbot = gr.Chatbot(label="Chat History", type="messages")
-    msg = gr.Textbox(label="Your message")
-    clear = gr.Button("Clear")
-
+    # Separate chatboxes for each model
     with gr.Row():
-        upvote = gr.Button("👍 Upvote")
-        downvote = gr.Button("👎 Downvote")
+        with gr.Column():
+            gr.Markdown("### 🧠 Model 1 Chat")
+            chatbot1 = gr.Chatbot(label=f"{models[0].split('/')[-1]} Chat History")
+        with gr.Column():
+            gr.Markdown("### 🧠 Model 2 Chat")
+            chatbot2 = gr.Chatbot(label=f"{models[1].split('/')[-1]} Chat History")
 
-    msg.submit(chat, [msg, chatbot, model1_dropdown, model2_dropdown], chatbot)
-    clear.click(lambda: [], None, chatbot, queue=False)
+    # Input textbox for user message
+    msg = gr.Textbox(label="💬 Your Message", placeholder="Type your message here...")
+
+    # Buttons for upvote, downvote, and clearing the chat
+    with gr.Row():
+        upvote = gr.Button("👍 Upvote Model 1")
+        downvote = gr.Button("👎 Downvote Model 2")
+        clear = gr.Button("🧹 Clear Chat")
+
+    # Textbox to display vote counts
+    vote_text = gr.Textbox(label="🏆 Vote Counts", value="Votes - 0, 0", interactive=False)
+
+    # Define interactions
+    msg.submit(
+        chat,
+        inputs=[msg, chatbot1, chatbot2, model1_dropdown, model2_dropdown],
+        outputs=[chatbot1, chatbot2]
+    )
+
+    upvote.click(
+        upvote_vote,
+        inputs=[model1_dropdown, model2_dropdown],
+        outputs=vote_text
+    )
+
+    downvote.click(
+        downvote_vote,
+        inputs=[model1_dropdown, model2_dropdown],
+        outputs=vote_text
+    )
+
+    clear.click(
+        clear_chat,
+        outputs=[chatbot1, chatbot2, vote_text]
+    )
+
+# Pre-load all models when the space starts
+load_all_models()
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()