# Example inference code
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load the model
tokenizer = AutoTokenizer.from_pretrained("exaler/aaa-2-sql-2")
model = AutoModelForCausalLM.from_pretrained("exaler/aaa-2-sql-2")

def generate_sql(instruction, input_text):
    # Format prompt
    prompt = f"<s>[INST] {instruction}\n\n{input_text} [/INST]"
    
    # Generate
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        inputs=inputs.input_ids, 
        max_new_tokens=512,
        temperature=0.0,
        do_sample=False
    )
    response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
    return response