from unsloth import FastLanguageModel
import torch
MODEL_NAME = "lakshitha722/querymind-nl2sql"
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = MODEL_NAME,
max_seq_length = 1024,
load_in_4bit = True,
dtype = None,
)
FastLanguageModel.for_inference(model)
PROMPT_TEMPLATE = """Below is an instruction that describes a task. Write a response that appropriately completes the request.
### Instruction:
Convert the following natural language question to a SQL query based on the given database schema. Return ONLY the SQL query, nothing else.
### Schema:
{schema}
### Question:
{question}
### Response:
"""
schema = "Database: company\nTables: employees (id, name, department, salary, hire_date)"
question = "What is the average salary by department?"
prompt = PROMPT_TEMPLATE.format(schema=schema, question=question)
inputs = tokenizer([prompt], return_tensors="pt").to("cuda")
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens = 150,
temperature = 0.1,
do_sample = False,
pad_token_id = tokenizer.eos_token_id,
)
input_length = inputs['input_ids'].shape[1]
sql = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True).strip()
print("Generated SQL:", sql)