Module_2 / model_loader.py
Springboardmen's picture
Update model_loader.py
4ad42f2 verified
raw
history blame
518 Bytes
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"]