

from transformers import GPTNeoXForCausalLM, GPTNeoXTokenizerFast
model = GPTNeoXForCausalLM.from_pretrained("EleutherAI/gpt-neox-20b")
tokenizer = GPTNeoXTokenizerFast.from_pretrained("EleutherAI/gpt-neox-20b")
prompt = "GPTNeoX20B is a 20B-parameter autoregressive Transformer model developed by EleutherAI."
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
gen_tokens = model.generate(
     input_ids,
     do_sample=True,
     temperature=0.9,
     max_length=100,
 )
gen_text = tokenizer.batch_decode(gen_tokens)[0]