example code for running GPT-Neo

gpt_neo_generation.py

"""
Invode the conda environment gpt_neo_generation before running this file.
Specify the prompt in prompt.txt
"""

from transformers import GPTNeoForCausalLM, GPT2Tokenizer
import time

def main():
    start_time = time.time()
    model = GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")
    tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
    num_times = 1
    temperature = 0.9
    max_length = 2000

    f = open("prompt.txt", "r")
    prompt = f.read()
    prompt = prompt.strip()

    for i in range(num_times):
        input_ids = tokenizer(prompt, return_tensors="pt").input_ids
        gen_tokens = model.generate(input_ids, do_sample=True, temperature=temperature, max_length=max_length)
        gen_text = tokenizer.batch_decode(gen_tokens)[0]
        print(gen_text)
        print("="*30)
    time_spent = time.time() - start_time
    print(f'spent {time_spent:.2f} seconds.')

if __name__ == "__main__":
    main()