Views
No views yet
b-mc2/sql-create-context dataset.1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
5
6model = AutoModelForCausalLM.from_pretrained("Artifact-io/toy-sql-28M").to(device)
7tokenizer = AutoTokenizer.from_pretrained("Artifact-io/toy-sql-28M")
8
9inputs = tokenizer([
10"""CREATE TABLE head (age INTEGER)
11How many heads of the departments are older than 56?
12"""
13 ],
14 return_tensors="pt",
15).to(device)
16
17outputs = model.generate(**inputs, max_new_tokens=200, do_sample=True, top_k=50, top_p=0.95)
18text = tokenizer.batch_decode(outputs[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)[0].split("---")[0]
19print(text)