Update model_loader.py

model_loader.py

@@ -1,13 +1,20 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-
-MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"
-
-def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        torch_dtype=torch.float16,
-        device_map="auto"
-    )
-    return tokenizer, model
+import requests
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"
+headers = {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
+
+def query_model(prompt):
+    payload = {
+        "inputs": prompt,
+        "parameters": {
+            "max_new_tokens": 800,
+            "temperature": 0.7
+        }
+    }
+
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.json()[0]["generated_text"]