File size: 718 Bytes
567c93e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
# 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
|