Update app.py

app.py

@@ -1,62 +1,85 @@
+import os
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import openai
-import torch
-
-# Load Llama model (GPU-optimized)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-llama_tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
-llama_model = AutoModelForCausalLM.from_pretrained(
-    "meta-llama/Llama-2-7b-chat-hf", 
-    device_map="auto"
-)
-
-# OpenAI GPT Model API Key (Replace with your API key)
-openai.api_key = "YOUR_OPENAI_API_KEY"
-
-# Function to query Llama
-def query_llama(prompt):
-    inputs = llama_tokenizer(prompt, return_tensors="pt", truncation=True, max_length=128).to(device)
-    outputs = llama_model.generate(inputs.input_ids, max_length=150)
-    response = llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
-
-# Function to query GPT
-def query_gpt(prompt):
-    response = openai.Completion.create(
-        engine="text-davinci-003",
-        prompt=prompt,
-        max_tokens=150
-    )
-    return response['choices'][0]['text'].strip()
+from huggingface_hub import login
+from huggingface_hub import InferenceClient
+import spaces
+
+# Authenticate with Hugging Face API
+api_key = os.getenv("LLAMA")
+login(api_key)
+
+# Initialize clients for different models
+llama_client = InferenceClient("meta-llama/Llama-3.1-70B-Instruct")
+gpt_client = InferenceClient("openai/gpt-4")  # Example: Replace with your OpenAI GPT model
+
+# Define the response function
+@spaces.GPU
+def respond(
+    message,
+    history: list[dict],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+    selected_models,
+):
+    # Prepare input messages
+    messages = [{"role": "system", "content": system_message}] + history
+    messages.append({"role": "user", "content": message})
 
-# Function to compare models
-def compare_models(prompt, models):
+    # Collect responses from selected models
     responses = {}
-    if "Llama" in models:
-        responses["Llama"] = query_llama(prompt)
-    if "GPT" in models:
-        responses["GPT"] = query_gpt(prompt)
+
+    if "Llama" in selected_models:
+        llama_response = ""
+        for token in llama_client.chat_completion(
+            messages, max_tokens=max_tokens, stream=True, temperature=temperature, top_p=top_p
+        ):
+            delta = token.choices[0].delta.content
+            llama_response += delta
+        responses["Llama"] = llama_response
+
+    if "GPT" in selected_models:
+        gpt_response = ""
+        for token in gpt_client.chat_completion(
+            messages, max_tokens=max_tokens, stream=True, temperature=temperature, top_p=top_p
+        ):
+            delta = token.choices[0].delta.content
+            gpt_response += delta
+        responses["GPT"] = gpt_response
+
     return responses
 
-# Gradio Interface
-def gradio_app():
-    with gr.Blocks() as app:
-        gr.Markdown("# AI Model Comparison Tool 🚀")
-        with gr.Row():
-            prompt_input = gr.Textbox(label="Enter your prompt", placeholder="Ask something...")
-        with gr.Row():
-            model_selector = gr.CheckboxGroup(
-                ["Llama", "GPT"], 
-                label="Select Models to Compare", 
-                value=["Llama", "GPT"]
-            )
-        with gr.Row():
-            output_boxes = gr.JSON(label="Model Responses")
-        with gr.Row():
-            compare_button = gr.Button("Compare Models")
-        compare_button.click(compare_models, inputs=[prompt_input, model_selector], outputs=[output_boxes])
-    return app
+# Build the Gradio app
+def create_demo():
+    return gr.Blocks().add(
+        gr.Markdown("# AI Model Comparison Tool 🌟"),
+        gr.ChatInterface(
+            respond,
+            type="messages",
+            additional_inputs=[
+                gr.Textbox(
+                    value="You are a helpful assistant providing answers for technical and customer support queries.",
+                    label="System message"
+                ),
+                gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+                gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+                gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.95,
+                    step=0.05,
+                    label="Top-p (nucleus sampling)"
+                ),
+                gr.CheckboxGroup(
+                    ["Llama", "GPT"],
+                    label="Select models to compare",
+                    value=["Llama"]
+                ),
+            ],
+        ),
+    )
 
 if __name__ == "__main__":
-    gradio_app().launch()
+    demo = create_demo()
+    demo.launch()