Update app.py

app.py

@@ -1,18 +1,26 @@
-from transformers import RobertaTokenizer, RobertaForMaskedLM, pipeline, GPT2Tokenizer, GPT2LMHeadModel
+from transformers import RobertaTokenizer, RobertaForMaskedLM, GPT2Tokenizer, AutoTokenizer, GPTJForCausalLM
 import torch
 import wikipedia
 import re
 import random
 import nltk
 import syllables
+from aitextgen import aitextgen
 nltk.download('cmudict')
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 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)
+
+# Without any parameters, aitextgen() will download, cache, and load the 124M GPT-2 "small" model
+gpt2 = aitextgen()
+
+gptj_tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")
+gptj_model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", revision="float16", low_cpu_mem_usage=True)
+gptj_model.to(device)
 
 frequent_words = set()
 
@@ -48,6 +56,7 @@ def filter_rhymes(word):
 
 def remove_punctuation(text):
     text = re.sub(r'[^\w\s]', '', text)
+    text = text.replace("\n", " ")
     return text
 
 def get_rhymes(inp, level):
@@ -100,19 +109,13 @@ def get_prediction(sent):
         
     return best_guess
     
-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:]
+    line = gpt2.generate_one(prompt=prompt + ".", max_length=inputs_len + 7)[len(prompt)+2:]
     return line
 
 def get_rhyming_line(prompt, rhyming_word, inputs_len):
-    gpt2_sentence = ai.generate_one(prompt=prompt + ".", max_length=inputs_len + 4)[len(prompt)+2:]
+    gpt2_sentence = gpt2.generate_one(prompt=prompt + ".", max_length=inputs_len + 4)[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
     print(f"Original Sentence: {sentence}")
@@ -128,12 +131,20 @@ def get_rhyming_line(prompt, rhyming_word, inputs_len):
     final_sentence = gpt2_sentence + predicted_blanks + " " + rhyming_word
     print(f"Final Sentence: {final_sentence}")
     return final_sentence
-    
-def generate(topic):
 
-    limericks = []
+def gptj_summary(topic):
+    prompt = f"Here is some information about {topic}"
+    input_ids = gptj_tokenizer(prompt, return_tensors="pt").input_ids.to(device)
 
-    topic_summary = remove_punctuation(wikipedia.summary(topic))
+    generated_ids = gptj_model.generate(input_ids, do_sample=True, temperature=0.9, max_length=200)
+    generated_text = gptj_tokenizer.decode(generated_ids[0])
+    return generated_text
+    
+def generate(topic, wiki=True):
+    if wiki:
+        topic_summary = remove_punctuation(wikipedia.summary(topic))
+    else:
+        topic_summary = remove_punctuation(gptj_summary(topic))
     word_list = topic_summary.split()
     topic_summary_len = len(topic_summary)
     no_of_words = len(word_list)
@@ -143,77 +154,82 @@ def generate(topic):
     print(f"No of Words in Summary: {no_of_words}")
     print(f"Length of Input IDs: {inputs_len}")         
 
-    for i in range(1):
-        print(f"\nGenerating limerick {i+1}")
-        rhyming_words_125 = []
-        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]
-        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 = []
-        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]
-        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]
-        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)
+    rhyming_words_125 = []
+    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]
+    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 = []
+    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]
+    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]
+    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"
 
     print("\n")
-    output = f"Generated {len(limericks)} limericks: \n"
+    print(limerick)
+
+    return limerick
 
-    print(f"Generated {len(limericks)} limericks: \n")
-    for limerick in limericks:
-        print(limerick)
-        output += "\n" + limerick
+def compare_summaries(topic):
+    wiki_limerick = generate(topic, wiki=True)
+    gptj_limerick = generate(topic, wiki=False)
+
+    output = f"Limerick with Wikipedia summary of topic as prompt: \n"
+    output += wiki_limerick + "\n"
+    output += f"Limerick with GPT-J summary of topic as prompt: \n"
+    output += gptj_limerick
 
     return output
-    
+      
 import gradio as gr
 
-interface = gr.Interface(fn=generate, inputs="text", outputs="text")
+interface = gr.Interface(
+    fn=compare_summaries, 
+    inputs="text", 
+    outputs="text")
 interface.launch(debug=True)