Update model_api.py

model_api.py

@@ -1,11 +1,13 @@
 import requests
-import streamlit as st
+import os
 import time
 
 API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"
 
+HF_TOKEN = os.getenv("HF_TOKEN")
+
 headers = {
-    "Authorization": f"Bearer {st.secrets['HF_TOKEN']}"
+    "Authorization": f"Bearer {HF_TOKEN}"
 }
 
 def query_model(prompt):
@@ -17,15 +19,14 @@ def query_model(prompt):
         }
     }
 
-    response = requests.post(API_URL, headers=headers, json=payload)
-    result = response.json()
-
-    # If model is loading
-    if isinstance(result, dict) and "error" in result:
-        return f"⚠ Model Error: {result['error']}"
+    while True:
+        response = requests.post(API_URL, headers=headers, json=payload)
+        result = response.json()
 
-    # If normal success response
-    if isinstance(result, list) and len(result) > 0:
-        return result[0]["generated_text"]
+        if isinstance(result, list):
+            return result[0]["generated_text"]
 
-    return "Unexpected response from model."
+        if "estimated_time" in result:
+            time.sleep(result["estimated_time"])
+        else:
+            return f"Error: {result}"