Update model_api.py

model_api.py

@@ -1,37 +1,25 @@
-import requests
+from huggingface_hub import InferenceClient
 import os
-import time
-
-API_URL = "https://router.huggingface.co/hf-inference/models/mistralai/Mistral-7B-Instruct-v0.2"
-
-HF_TOKEN = os.getenv("HF_TOKEN")
-
-headers = {
-    "Authorization": f"Bearer {HF_TOKEN}"
-}
 
 def query_model(prompt):
-    payload = {
-        "inputs": prompt,
-        "parameters": {
-            "max_new_tokens": 500,
-            "temperature": 0.7
-        }
-    }
-
-    response = requests.post(API_URL, headers=headers, json=payload)
+    try:
+        HF_TOKEN = os.getenv("HF_TOKEN")
 
-    # ✅ Check status code first
-    if response.status_code != 200:
-        return f"HTTP Error {response.status_code}: {response.text}"
+        if not HF_TOKEN:
+            return "⚠ HF_TOKEN not found."
 
+        client = InferenceClient(
+            model="mistralai/Mistral-7B-Instruct-v0.2",
+            token=HF_TOKEN
+        )
 
-    try:
-        result = response.json()
-    except Exception:
-        return f"Invalid response received: {response.text}"
+        response = client.text_generation(
+            prompt,
+            max_new_tokens=300,
+            temperature=0.7
+        )
 
-    if isinstance(result, list):
-        return result[0]["generated_text"]
+        return response
 
-    return f"Unexpected response: {result}"
+    except Exception as e:
+        return f"Error: {str(e)}"