from transformers import AutoModelForCausalLM, AutoTokenizerimport 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