Update app.py

app.py

@@ -1,41 +1,19 @@
-from transformers import RobertaTokenizer, RobertaForMaskedLM, GPT2Tokenizer, GPTNeoForCausalLM
+from transformers import RobertaTokenizer, RobertaForMaskedLM, GPT2Tokenizer
 import torch
 import wikipedia
 import re
 import random
 import nltk
-import syllables
 from aitextgen import aitextgen
 nltk.download('cmudict')
 
-masked_tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
-masked_model = RobertaForMaskedLM.from_pretrained('roberta-base')
+roberta_tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
+roberta_model = RobertaForMaskedLM.from_pretrained('roberta-base')
 
-causal_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-gptneo_tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
-gptneo_model = GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")
-
-# Without any parameters, aitextgen() will download, cache, and load the 124M GPT-2 "small" model
-gpt2 = aitextgen()
+gpt2_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+gpt2_model = aitextgen(tf_gpt2="355M")
 
 frequent_words = set()
-
-def set_seed(seed: int):
-    """
-    Helper function for reproducible behavior to set the seed in ``random``, ``numpy``, ``torch`` and/or ``tf`` (if
-    installed).
-
-    Args:
-        seed (:obj:`int`): The seed to set.
-    """
-    #random.seed(seed)
-    #np.random.seed(seed)
-    #if is_torch_available():
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-        # ^^ safe to call this function even if cuda is not available
-    #if is_tf_available():
-        #tf.random.set_seed(seed)
         
 with open("wordFrequency.txt", 'r') as f:
     line = f.readline()
@@ -69,19 +47,17 @@ def get_rhymes(inp, level):
     return filtered_rhymes
 
 def get_inputs_length(input):
-    input_ids = causal_tokenizer(input)['input_ids']
+    input_ids = gpt2_tokenizer(input)['input_ids']
     return len(input_ids)
-       
-set_seed(0)
     
 def get_prediction(sent):
     
-    token_ids = masked_tokenizer.encode(sent, return_tensors='pt')
-    masked_position = (token_ids.squeeze() == masked_tokenizer.mask_token_id).nonzero()
+    token_ids = roberta_tokenizer.encode(sent, return_tensors='pt')
+    masked_position = (token_ids.squeeze() == roberta_tokenizer.mask_token_id).nonzero()
     masked_pos = [mask.item() for mask in masked_position ]
 
     with torch.no_grad():
-        output = masked_model(token_ids)
+        output = roberta_model(token_ids)
 
     last_hidden_state = output[0].squeeze()
 
@@ -92,10 +68,9 @@ def get_prediction(sent):
             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()
+                word = roberta_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}")
     
@@ -104,13 +79,13 @@ def get_prediction(sent):
         best_guess = best_guess+" "+j[0]
         
     return best_guess
-    
+
 def get_line(prompt, inputs_len):
-    line = gpt2.generate_one(prompt=prompt + ".", max_length=inputs_len + 7)[len(prompt)+2:]
+    line = gpt2_model.generate_one(prompt=prompt + ".", max_length=inputs_len + 7, min_length=4)[len(prompt)+2:]
     return line
 
 def get_rhyming_line(prompt, rhyming_word, inputs_len):
-    gpt2_sentence = gpt2.generate_one(prompt=prompt + ".", max_length=inputs_len + 4)[len(prompt)+2:]
+    gpt2_sentence = gpt2_model.generate_one(prompt=prompt + ".", max_length=inputs_len + 4, min_length=2)[len(prompt)+2:]
     gpt2_sentence = gpt2_sentence.replace("\n", "")
     print(f"\nGetting rhyming line starting with '{gpt2_sentence}' and ending with rhyming word '{rhyming_word}'")
     sentence = gpt2_sentence + " ___ ___ ___ " + rhyming_word
@@ -128,17 +103,15 @@ def get_rhyming_line(prompt, rhyming_word, inputs_len):
     print(f"Final Sentence: {final_sentence}")
     return final_sentence
 
-def gptneo_summary(topic):
-    input_ids = gptneo_tokenizer(f"Here is some information about {topic}", return_tensors="pt").input_ids
-    gen_tokens = gptneo_model.generate(input_ids, do_sample=True, temperature=0.9, max_length=200)
-    generated_text = gptneo_tokenizer.decode(gen_tokens[0])
-    return generated_text 
-    
+def gpt2_summary(topic):
+    return gpt2_model.generate_one(prompt=f"Here is some information about {topic}", top_k=100, top_p=0.95)
+
 def generate(topic, wiki=True):
     if wiki:
         topic_summary = remove_punctuation(wikipedia.summary(topic))
     else:
-        topic_summary = remove_punctuation(gptneo_summary(topic))
+        topic_summary = remove_punctuation(gpt2_summary(topic))
+
     word_list = topic_summary.split()
     topic_summary_len = len(topic_summary)
     no_of_words = len(word_list)
@@ -160,7 +133,8 @@ def generate(topic, wiki=True):
                 print(f"Rhyming words for '{end_word}' are {rhyming_words_125}")
                 limerick = first_line + "\n"
 
-    rhyming_word = rhyming_words_125[0]
+    rhyming_word = random.choice(rhyming_words_125)
+    rhyming_words_125.remove(rhyming_word)
     prompt = topic_summary + " " + first_line
     inputs_len = get_inputs_length(prompt)
     print(f"Prompt: {prompt}")
@@ -186,7 +160,8 @@ def generate(topic, wiki=True):
             if valid_rhyme and len(rhyming_words_34) > 1:
                 limerick += third_line + "\n"
 
-    rhyming_word = rhyming_words_34[0]
+    rhyming_word = random.choice(rhyming_words_34)
+    rhyming_words_34.remove(rhyming_word)
     prompt = prompt + " " + third_line
     inputs_len = get_inputs_length(prompt)
     print(f"Prompt: {prompt}")
@@ -195,7 +170,8 @@ def generate(topic, wiki=True):
     print(f"\nFourth Line: {fourth_line}")
     limerick += fourth_line + "\n"
 
-    rhyming_word = rhyming_words_125[1]
+    rhyming_word = random.choice(rhyming_words_125)
+    rhyming_words_125.remove(rhyming_word)
     prompt = prompt + " " + fourth_line
     inputs_len = get_inputs_length(prompt)
     print(f"Prompt: {prompt}")
@@ -210,20 +186,23 @@ def generate(topic, wiki=True):
     return limerick
 
 def compare_summaries(topic):
-    wiki_limerick = generate(topic, wiki=True)
-    gptneo_limerick = generate(topic, wiki=False)
+    wiki_limerick = generate(topic)
+    gpt2_limerick = generate(topic, wiki=False)
+
+    output1 = f"Limerick with Wikipedia summary of topic as prompt: \n"
+    output1 += wiki_limerick + "\n"
+    output2 = f"Limerick with GPT-2 summary of topic as prompt: \n"
+    output2 += gpt2_limerick
 
-    output = f"Limerick with Wikipedia summary of topic as prompt: \n"
-    output += wiki_limerick + "\n"
-    output += f"Limerick with GPT Neo summary of topic as prompt: \n"
-    output += gptneo_limerick
+    return output1, output2
 
-    return output
-      
 import gradio as gr
 
 interface = gr.Interface(
     fn=compare_summaries, 
     inputs="text", 
-    outputs="text")
+    outputs=["text", "text"],
+    title="Text-generation with rhyme and rhythm",
+    layout="horizontal",
+    theme="peach")
 interface.launch(debug=True)