import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
model_id = "Qwen/Qwen2.5-0.5B-Instruct"
adapter_id = "azeemazam/Qwen2.5-0.5B-SQL"
tokenizer = AutoTokenizer.from_pretrained(adapter_id)
base_model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map='auto')
model = PeftModel.from_pretrained(base_model, adapter_id)
def generate_sql(schema, question):
messages = [
{"role": "user", "content": f"Generate SQL.\\n\\nDatabase Schema:\\n{schema}\\n\\nQuestion:\\n{question}"}
]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
outputs = model.generate(**inputs, max_new_tokens=150)
return tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
schema = "CREATE TABLE employees (id INT, name TEXT, salary INT)"
question = "Who earns more than 50000?"
print(generate_sql(schema, question))