AuricErgeson/Antelope-textTosql

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "AuricErgeson/Antelope-textTosql"

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True
)

def generate_sql(question, db=""):
    prompt = (
        f"### Task: Convert question to SQL. Use only what the question asks. Simple questions need simple SQL.\n"
        f"### Database: {db}\n"
        f"### Question: {question}\n"
        f"### SQL:"
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=128,
            temperature=0.1,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    decoded = tokenizer.decode(out[0], skip_special_tokens=True)
    return decoded.split("### SQL:")[-1].strip().split("\n")[0]

print(generate_sql("How many employees are in each department?", db="company"))
# → SELECT department, COUNT(*) FROM employees GROUP BY department

### Task: Convert question to SQL. Use only what the question asks. Simple questions need simple SQL.
### Database: {your_database_name}
### Question: {your_question}
### SQL:

import requests

API_URL = "https://api-inference.huggingface.co/models/AuricErgeson/Antelope-textTosql"
headers = {"Authorization": "Bearer YOUR_HF_TOKEN"}

response = requests.post(API_URL, headers=headers, json={
    "inputs": "### Task: Convert question to SQL. Use only what the question asks. Simple questions need simple SQL.\n### Database: company\n### Question: How many employees?\n### SQL:"
})
print(response.json())