augment_generate.py

from transformers import pipeline
generator = pipeline('text-generation', model='gpt2')

input_text = "I went to see a movie in the theater"
input_length = len(input_text.split())
num_new_words = 5
output_length = input_length + num_new_words
gpt_output = generator(input_text, max_length=output_length, num_return_sequences=5)
augmented_text = gpt_output[0]['generated_text']
print("Augmented text->",augmented_text)
#I went to see a movie in the theater, and the director was