Artifact-io/toy-sql-28M

Toy model finetuned on the b-mc2/sql-create-context dataset.

Sample Code

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = AutoModelForCausalLM.from_pretrained("Artifact-io/toy-sql-28M").to(device)
tokenizer = AutoTokenizer.from_pretrained("Artifact-io/toy-sql-28M")

inputs = tokenizer([
"""CREATE TABLE head (age INTEGER)
How many heads of the departments are older than 56?
"""
  ],
  return_tensors="pt",
).to(device)

outputs = model.generate(**inputs, max_new_tokens=200, do_sample=True, top_k=50, top_p=0.95)
text = tokenizer.batch_decode(outputs[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)[0].split("---")[0]
print(text)