Update app.py

app.py

@@ -1,73 +1,63 @@
 import gradio as gr
 import os
 from openai import OpenAI
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Set your API keys as environment variables or replace os.getenv with your actual keys
 DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 
-# Initialize OpenAI client (for openai>=1.0.0)
-client = OpenAI(api_key=OPENAI_API_KEY)
-
-# Load DeepSeek model
-deepseek_model_id = "deepseek-ai/deepseek-llm-7b-chat"
-tokenizer = AutoTokenizer.from_pretrained(deepseek_model_id)
-deepseek_model = AutoModelForCausalLM.from_pretrained(
-    deepseek_model_id,
-    torch_dtype=torch.float16,
-    device_map="auto"
-)
+# Initialize OpenAI clients
+openai_client = OpenAI(api_key=OPENAI_API_KEY)
+deepseek_client = OpenAI(api_key=DEEPSEEK_API_KEY, base_url="https://api.deepseek.com")
 
 def generate_response(prompt, model_provider, temperature, top_p, max_tokens, repetition_penalty):
     if model_provider == "DeepSeek":
-        inputs = tokenizer(prompt, return_tensors="pt").to(deepseek_model.device)
-        outputs = deepseek_model.generate(
-            **inputs,
-            do_sample=True,
-            temperature=temperature,
-            top_p=top_p,
-            max_new_tokens=max_tokens,
-            repetition_penalty=repetition_penalty
-        )
-        return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
+        try:
+            response = deepseek_client.chat.completions.create(
+                model="deepseek-chat",  # or "deepseek-reasoner" for R1 model
+                messages=[{"role": "user", "content": prompt}],
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                presence_penalty=repetition_penalty,
+                stream=False
+            )
+            return response.choices[0].message.content.strip()
+        except Exception as e:
+            return f"DeepSeek API Error: {str(e)}"
     elif model_provider == "OpenAI":
         try:
-            response = client.chat.completions.create(
-                model="gpt-3.5-turbo",  # or "gpt-4" if you have access
+            response = openai_client.chat.completions.create(
+                model="gpt-3.5-turbo",  # or another model of your choice
                 messages=[{"role": "user", "content": prompt}],
                 temperature=temperature,
                 top_p=top_p,
                 max_tokens=max_tokens,
-                presence_penalty=repetition_penalty
+                presence_penalty=repetition_penalty,
+                stream=False
             )
             return response.choices[0].message.content.strip()
         except Exception as e:
             return f"OpenAI API Error: {str(e)}"
-
     else:
         return "Invalid model provider selected."
 
 with gr.Blocks() as demo:
     gr.Markdown("## 🔍 LLM Chat Interface")
-
     with gr.Row():
         model_provider = gr.Dropdown(
             choices=["DeepSeek", "OpenAI"],
             value="DeepSeek",
             label="Select Model Provider"
         )
-
     prompt = gr.Textbox(label="Enter your prompt", lines=4, placeholder="Type your message here...")
-
     with gr.Accordion("Advanced Settings", open=False):
         temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperature")
         top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
         max_tokens = gr.Slider(32, 2048, value=512, step=32, label="Max New Tokens")
         repetition_penalty = gr.Slider(1.0, 2.0, value=1.1, step=0.1, label="Repetition Penalty")
-
     output = gr.Textbox(label="Response")
     submit = gr.Button("Generate")