from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
base = "Qwen/Qwen2.5-Coder-1.5B-Instruct"
tok = AutoTokenizer.from_pretrained(base)
model = AutoModelForCausalLM.from_pretrained(base, device_map="auto")
model = PeftModel.from_pretrained(model, "Anurich/slmsql-qwen2.5-coder-1.5b")
messages = [
{"role": "system", "content": "You are an expert data analyst who translates "
"questions into correct SQLite queries. Think step by step, then output the "
"final query in a ```sql block. Wrap identifiers with spaces in backticks."},
{"role": "user", "content": "Given the following database schema:\n\n"
"CREATE TABLE employees (id INTEGER, name TEXT, salary INTEGER);\n\n"
"Question: Who are the 3 highest-paid employees?\n\n"
"Write a single SQLite query that answers the question."},
]
prompt = tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
out = model.generate(**tok(prompt, return_tensors="pt").to(model.device), max_new_tokens=512)
print(tok.decode(out[0], skip_special_tokens=True))