Update app.py

app.py

@@ -6,6 +6,12 @@ from threading import Thread
 import re
 import uuid
 
+client = OpenAI(
+    base_url="https://a7g1ajqixo23revq.us-east-1.aws.endpoints.huggingface.cloud/v1/",
+    api_key="hf_XXXXX"
+)
+
+
 # Load model and tokenizer
 our_model_path = "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
@@ -24,8 +30,8 @@ conversations = {}
 def generate_conversation_id():
     return str(uuid.uuid4())[:8]
 
-@spaces.GPU(duration=60)
-def generate_response(user_message, max_tokens, temperature, top_p, history_state):
+#@spaces.GPU(duration=60)
+'''def generate_response(user_message, max_tokens, temperature, top_p, history_state):
     if not user_message.strip():
         return history_state, history_state
 
@@ -85,8 +91,57 @@ def generate_response(user_message, max_tokens, temperature, top_p, history_stat
         pass
 
     yield new_history, new_history
+'''
+
+@spaces.GPU(duration=60)
+def generate_response(user_message, max_tokens, temperature, top_p, history_state):
+    if not user_message.strip():
+        return history_state, history_state
+
+    start_tag = "<|im_start|>"
+    sep_tag = "<|im_sep|>"
+    end_tag = "<|im_end|>"
+
+    system_message = "Your role as an assistant..."
+    messages = [{"role": "system", "content": system_message}]
+    for message in history_state:
+        messages.append({"role": message["role"], "content": message["content"]})
+    messages.append({"role": "user", "content": user_message})
 
+    try:
+        response = client.chat.completions.create(
+            model="tgi",
+            messages=messages,
+            max_tokens=int(max_tokens),
+            temperature=temperature,
+            top_p=top_p,
+            stream=True,
+        )
+    except Exception as e:
+        yield history_state + [{"role": "user", "content": user_message},
+                               {"role": "assistant", "content": "⚠️ Generation failed."}], history_state
+        return
+
+    assistant_response = ""
+    new_history = history_state + [
+        {"role": "user", "content": user_message},
+        {"role": "assistant", "content": ""}
+    ]
+
+    try:
+        for chunk in response:
+            if not chunk.choices or not chunk.choices[0].delta or not chunk.choices[0].delta.content:
+                continue
+            token = chunk.choices[0].delta.content
+            assistant_response += token
+            new_history[-1]["content"] = assistant_response.strip()
+            yield new_history, new_history
+    except Exception:
+        pass
+
+    yield new_history, new_history
 
+    
 example_messages = {
     "JEE Main 2025 Combinatorics": "From all the English alphabets, five letters are chosen and are arranged in alphabetical order. The total number of ways, in which the middle letter is 'M', is?",
     "JEE Main 2025 Coordinate Geometry": "A circle \\(C\\) of radius 2 lies in the second quadrant and touches both the coordinate axes. Let \\(r\\) be the radius of a circle that has centre at the point \\((2, 5)\\) and intersects the circle \\(C\\) at exactly two points. If the set of all possible values of \\(r\\) is the interval \\((\\alpha, \\beta)\\), then \\(3\\beta - 2\\alpha\\) is?",