Skip to content

Instantly share code, notes, and snippets.

@suryavanshi

suryavanshi/augment_generate.py

Created August 13, 2021 10:52
Show Gist options
  • Save suryavanshi/5e3758bf957c5afa4e8f35a31c568d06 to your computer and use it in GitHub Desktop.
Save suryavanshi/5e3758bf957c5afa4e8f35a31c568d06 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment