import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
base_model_name = "Qwen/Qwen2.5-1.5B-Instruct"
adapter_repo = "Amr-Kh-2004/Qwen2.5-1.5B-Instruct-text2sql-lora"
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
)
tokenizer = AutoTokenizer.from_pretrained(adapter_repo)
base_model = AutoModelForCausalLM.from_pretrained(
base_model_name,
device_map="auto",
quantization_config=quantization_config,
torch_dtype=torch.float16,
)
model = PeftModel.from_pretrained(base_model, adapter_repo)
model.eval()
question = "List all employees hired after 2020."
db_schema = "CREATE TABLE employees (id INT, name VARCHAR(50), hire_date DATE);"
chat = [
{"role": "system", "content": "You are an SQL query generator given a question and a database schema, generate an SQL query."},
{"role": "user", "content": f"Question: {question}\nDatabase Schema: {db_schema}\nSQL Query:"},
]
prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
input_len = inputs["input_ids"].shape[1]
output = model.generate(**inputs, max_new_tokens=256, do_sample=False)
print(tokenizer.decode(output[0][input_len:], skip_special_tokens=True))