| ) | |
| 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] | |
| or in float16 precision: | |
| thon | |
| from transformers import GPTJForCausalLM, AutoTokenizer | |
| import torch | |
| device = "cuda" | |
| model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", torch_dtype=torch.float16).to(device) | |
| tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B") | |
| prompt = ( | |
| "In a shocking finding, scientists discovered a herd of unicorns living in a remote, " | |
| "previously unexplored valley, in the Andes Mountains. |