| # 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 | |