Example
gen.py
from transformers import GPTNeoForCausalLM, AutoTokenizer
import torch
import sys
model_name = sys.argv[1]
model = GPTNeoForCausalLM.from_pretrained(model_name).to("cuda")
tokenizer = AutoTokenizer.from_pretrained(model_name)
def generate(model, text, temperature=0.9, min_length=256, max_length=256, no_grad=True, use_cache=False, do_sample=True, match_mesh_tf=False, **kwargs):
ids = tokenizer(text, return_tensors="pt").input_ids.to("cuda")
if no_grad:
with torch.no_grad():
gen_tokens = model.generate(
ids,
do_sample=do_sample,
min_length=min_length,
max_length=max_length,
temperature=temperature,
use_cache=use_cache,
**kwargs
)
gen_text = tokenizer.batch_decode(gen_tokens)[0]
print(gen_text)
python gen.py spamtontalk_gpt_neo_xl_v9
>>> text = """Talk (anything):
Example dialogue"""
>>> generate(model, text, temperature=0.92)
