infer.py

import os
import requests

API_URL = (
    "/static-proxy?url=https%3A%2F%2Fapi-inference.huggingface.co%2Fmodels%2Fmistralai%2FMistral-7B-Instruct-v0.2%26quot%3B%3C%2Fspan%3E%3C!-- HTML_TAG_END -->
)


def query(question, hub_token: str):
    payload = {
        "inputs": question,
        "parameters": {
            "wait_for_model": True,
            "return_full_text": False,
        },
    }
    headers = {"Authorization": f"Bearer {hub_token}"}
    response = requests.post(API_URL, headers=headers, json=payload)
    try:
        return response.json()[0]["generated_text"]
    except Exception:
        return "Error occurred while querying the model."