Update app.py

app.py

@@ -1,13 +1,19 @@
-from transformers import RobertaTokenizer, RobertaForMaskedLM, pipeline, GPT2TokenizerFast
+from transformers import RobertaTokenizer, RobertaForMaskedLM, pipeline, GPT2Tokenizer, GPT2LMHeadModel
 import torch
 import wikipedia
 import re
 import random
 import nltk
 import syllables
-import gradio as gr
 nltk.download('cmudict')
 
+masked_tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
+masked_model = RobertaForMaskedLM.from_pretrained('roberta-base')
+
+causal_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+# add the EOS token as PAD token to avoid warnings
+causal_model = GPT2LMHeadModel.from_pretrained("gpt2", pad_token_id=causal_tokenizer.eos_token_id)
+
 frequent_words = set()
 
 def set_seed(seed: int):
@@ -58,36 +64,33 @@ def get_rhymes(inp, level):
     return filtered_rhymes
 
 def get_inputs_length(input):
-    gpt2_tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
-    input_ids = gpt2_tokenizer(input)['input_ids']
+    input_ids = causal_tokenizer(input)['input_ids']
     return len(input_ids)
-
-tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
-model = RobertaForMaskedLM.from_pretrained('roberta-base')
-text_generation = pipeline("text-generation")
+    
 set_seed(0)
     
 def get_prediction(sent):
     
-    token_ids = tokenizer.encode(sent, return_tensors='pt')
-    masked_position = (token_ids.squeeze() == tokenizer.mask_token_id).nonzero()
+    token_ids = masked_tokenizer.encode(sent, return_tensors='pt')
+    masked_position = (token_ids.squeeze() == masked_tokenizer.mask_token_id).nonzero()
     masked_pos = [mask.item() for mask in masked_position ]
 
     with torch.no_grad():
-        output = model(token_ids)
+        output = masked_model(token_ids)
 
     last_hidden_state = output[0].squeeze()
 
     list_of_list =[]
     for index,mask_index in enumerate(masked_pos):
         words = []
-        mask_hidden_state = last_hidden_state[mask_index]
-        idx = torch.topk(mask_hidden_state, k=5, dim=0)[1]
-        for i in idx:
-            word = tokenizer.decode(i.item()).strip()
-            if (remove_punctuation(word) != "") and (word != '</s>'):
-                words.append(word)
-        #words = [tokenizer.decode(i.item()).strip() for i in idx]
+        while not words:
+            mask_hidden_state = last_hidden_state[mask_index]
+            idx = torch.topk(mask_hidden_state, k=5, dim=0)[1]
+            for i in idx:
+                word = masked_tokenizer.decode(i.item()).strip()
+                if (remove_punctuation(word) != "") and (word != '</s>'):
+                    words.append(word)
+            #words = [masked_tokenizer.decode(i.item()).strip() for i in idx]
         list_of_list.append(words)
         print(f"Mask {index+1} Guesses: {words}")
     
@@ -97,18 +100,21 @@ def get_prediction(sent):
         
     return best_guess
     
-def get_line(topic_summary, starting_words, inputs_len):
-    starting_word = random.choice(starting_words)
-    line = starting_word + text_generation(topic_summary + " " + starting_word, max_length=inputs_len + 6, do_sample=True, return_full_text=False)[0]['generated_text']
+text_generation = pipeline("text-generation", model=causal_model, tokenizer=causal_tokenizer)
+from aitextgen import aitextgen
+
+# Without any parameters, aitextgen() will download, cache, and load the 124M GPT-2 "small" model
+ai = aitextgen()
+
+
+def get_line(prompt, inputs_len):
+    line = ai.generate_one(prompt=prompt + ".", max_length=inputs_len + 7)[len(prompt)+2:]
     return line
 
-def get_rhyming_line(topic_summary, starting_words, rhyming_word, inputs_len):
-    #gpt2_sentence = text_generation(topic_summary + " " + starting_words[i][j], max_length=no_of_words + 4, do_sample=False)[0]
-    starting_word = random.choice(starting_words)
-    print(f"\nGetting rhyming line with starting word '{starting_word}' and rhyming word '{rhyming_word}'")
-    gpt2_sentence = text_generation(topic_summary + " " + starting_word, max_length=inputs_len + 2, do_sample=True, return_full_text=False)[0]
-    #sentence = gpt2_sentence['generated_text'] + " ___ ___ ___ " + rhyming_words[i][j]
-    sentence = starting_word + gpt2_sentence['generated_text'] + " ___ ___ ___ " + rhyming_word
+def get_rhyming_line(prompt, rhyming_word, inputs_len):
+    gpt2_sentence = ai.generate_one(prompt=prompt + ".", max_length=inputs_len + 4)[len(prompt)+2:]
+    print(f"\nGetting rhyming line starting with '{gpt2_sentence}' and ending with rhyming word '{rhyming_word}'")
+    sentence = gpt2_sentence + " ___ ___ ___ " + rhyming_word
     print(f"Original Sentence: {sentence}")
     if sentence[-1] != ".":
         sentence = sentence.replace("___","<mask>") + "."
@@ -119,20 +125,15 @@ def get_rhyming_line(topic_summary, starting_words, rhyming_word, inputs_len):
  
     predicted_blanks = get_prediction(sentence)
     print(f"\nBest guess for fill in the blanks: {predicted_blanks}")
-    return starting_word + gpt2_sentence['generated_text'] + predicted_blanks + " " + rhyming_word
+    final_sentence = gpt2_sentence + predicted_blanks + " " + rhyming_word
+    print(f"Final Sentence: {final_sentence}")
+    return final_sentence
     
-from transformers import pipeline
-
 def generate(topic):
-    text_generation = pipeline("text-generation")
 
     limericks = []
 
-    #topic = input("Please enter a topic: ")
-    ## topic_summary = remove_punctuation(wikipedia.summary(topic))
-    topic_summary = topic
-    # if len(topic_summary) > 2000:
-    #   topic_summary = topic_summary[:2000]
+    topic_summary = remove_punctuation(wikipedia.summary(topic))
     word_list = topic_summary.split()
     topic_summary_len = len(topic_summary)
     no_of_words = len(word_list)
@@ -140,59 +141,64 @@ def generate(topic):
     print(f"Topic Summary: {topic_summary}")
     print(f"Topic Summary Length: {topic_summary_len}")
     print(f"No of Words in Summary: {no_of_words}")
-    print(f"Length of Input IDs: {inputs_len}")
-
-    starting_words = ["That", "Had", "Not", "But", "With", "I", "Because", "There", "Who", "She", "He", "To", "Whose", "In", "And", "When", "Or", "So", "The", "Of", "Every", "Whom"]
-
-    # starting_words = [["That", "Had", "Not", "But", "That"], 
-    #                   ["There", "Who", "She", "Tormenting", "Til"],
-    #                   ["Relentless", "This", "First", "and", "then"],
-    #                   ["There", "Who", "That", "To", "She"],
-    #                   ["There", "Who", "Two", "Four", "Have"]]
+    print(f"Length of Input IDs: {inputs_len}")         
 
-    # rhyming_words = [["told", "bold", "woodchuck", "truck", "road"], 
-    #                  ["Nice", "grease", "house", "spouse", "peace"],
-    #                  ["deadlines", "lines", "edits", "credits", "wine"],
-    #                  ["Lynn", "thin", "essayed", "lemonade", "in"],
-    #                  ["beard", "feared", "hen", "wren", "beard"]]                 
-
-    for i in range(5):
+    for i in range(1):
         print(f"\nGenerating limerick {i+1}")
         rhyming_words_125 = []
-        while len(rhyming_words_125) < 3 or valid_rhyme == False:
-            first_line = get_line(topic_summary, starting_words, inputs_len)
-            #rhyming_words = pronouncing.rhymes(first_line.split()[-1])
-            end_word = remove_punctuation(first_line.split()[-1])
-            valid_rhyme = filter_rhymes(end_word)
-            if valid_rhyme:
-                print(f"\nFirst Line: {first_line}")
-                rhyming_words_125 = list(get_rhymes(end_word, 3))
-                print(f"Rhyming words for '{end_word}' are {rhyming_words_125}")
-                limerick = first_line + "\n"
+        while len(rhyming_words_125) < 3 or valid_rhyme == False or len(first_line) == 0:
+            first_line = get_line(topic_summary, inputs_len)
+            if first_line:
+                end_word = remove_punctuation(first_line.split()[-1])
+                valid_rhyme = filter_rhymes(end_word)
+                if valid_rhyme:
+                    print(f"\nFirst Line: {first_line}")
+                    rhyming_words_125 = list(get_rhymes(end_word, 3))
+                    print(f"Rhyming words for '{end_word}' are {rhyming_words_125}")
+                    limerick = first_line + "\n"
 
         rhyming_word = rhyming_words_125[0]
-        second_line = get_rhyming_line(topic_summary, starting_words, rhyming_word, inputs_len)
+        prompt = topic_summary + " " + first_line
+        inputs_len = get_inputs_length(prompt)
+        print(f"Prompt: {prompt}")
+        print(f"Length of prompt: {inputs_len}")
+        second_line = get_rhyming_line(prompt, rhyming_word, inputs_len)
+        print(f"\nSecond Line: {second_line}")
         limerick += second_line + "\n"
 
         rhyming_words_34 = []
-        while len(rhyming_words_34) < 2 or valid_rhyme == False:
-            third_line = get_line(topic_summary, starting_words, inputs_len)
-            print(f"\nThird Line: {third_line}")
-            #rhyming_words = pronouncing.rhymes(first_line.split()[-1])
-            end_word = remove_punctuation(third_line.split()[-1])
-            valid_rhyme = filter_rhymes(end_word)
-            print(f"Does '{end_word}'' have valid rhymes: {valid_rhyme}")
-            rhyming_words_34 = list(get_rhymes(end_word, 3))
-            print(f"Rhyming words for '{end_word}' are {rhyming_words_34}")
-            if valid_rhyme and len(rhyming_words_34) > 1:
-                limerick += third_line + "\n"
+        prompt = prompt + " " + second_line
+        inputs_len = get_inputs_length(prompt)
+        print(f"Prompt: {prompt}")
+        print(f"Length of prompt: {inputs_len}")
+        while len(rhyming_words_34) < 2 or valid_rhyme == False or len(third_line) == 0:
+            third_line = get_line(prompt, inputs_len)
+            if third_line:
+                print(f"\nThird Line: {third_line}")
+                end_word = remove_punctuation(third_line.split()[-1])
+                valid_rhyme = filter_rhymes(end_word)
+                print(f"Does '{end_word}'' have valid rhymes: {valid_rhyme}")
+                rhyming_words_34 = list(get_rhymes(end_word, 3))
+                print(f"Rhyming words for '{end_word}' are {rhyming_words_34}")
+                if valid_rhyme and len(rhyming_words_34) > 1:
+                    limerick += third_line + "\n"
 
         rhyming_word = rhyming_words_34[0]
-        fourth_line = get_rhyming_line(topic_summary, starting_words, rhyming_word, inputs_len)
+        prompt = prompt + " " + third_line
+        inputs_len = get_inputs_length(prompt)
+        print(f"Prompt: {prompt}")
+        print(f"Length of prompt: {inputs_len}")
+        fourth_line = get_rhyming_line(prompt, rhyming_word, inputs_len)
+        print(f"\nFourth Line: {fourth_line}")
         limerick += fourth_line + "\n"
 
         rhyming_word = rhyming_words_125[1]
-        fifth_line = get_rhyming_line(topic_summary, starting_words, rhyming_word, inputs_len)
+        prompt = prompt + " " + fourth_line
+        inputs_len = get_inputs_length(prompt)
+        print(f"Prompt: {prompt}")
+        print(f"Length of prompt: {inputs_len}")
+        fifth_line = get_rhyming_line(prompt, rhyming_word, inputs_len)
+        print(f"\nFifth Line: {fifth_line}")
         limerick += fifth_line + "\n"
 
         limericks.append(limerick)
@@ -203,9 +209,11 @@ def generate(topic):
     print(f"Generated {len(limericks)} limericks: \n")
     for limerick in limericks:
         print(limerick)
-        output += limerick
+        output += "\n" + limerick
 
     return output
+    
+import gradio as gr
 
 interface = gr.Interface(fn=generate, inputs="text", outputs="text")
 interface.launch(debug=True)