Update app.py

app.py

@@ -4,6 +4,7 @@ import wikipedia
 import re
 import random
 import nltk
+import gradio as gr
 from aitextgen import aitextgen
 nltk.download('cmudict')
 
@@ -49,7 +50,7 @@ def get_rhymes(inp, level):
 def get_inputs_length(input):
     input_ids = gpt2_tokenizer(input)['input_ids']
     return len(input_ids)
-	
+    
 def get_prediction(sent):
     
     token_ids = roberta_tokenizer.encode(sent, return_tensors='pt')
@@ -79,7 +80,7 @@ def get_prediction(sent):
         best_guess = best_guess+" "+j[0]
         
     return best_guess
-	
+
 def get_line(prompt, inputs_len):
     line = gpt2_model.generate_one(prompt=prompt + ".", max_length=inputs_len + 7, min_length=4)[len(prompt)+2:]
     return line
@@ -107,12 +108,16 @@ def get_rhyming_line(prompt, rhyming_word, inputs_len):
     return final_sentence
 
 def gpt2_summary(topic):
-    return gpt2_model.generate_one(prompt=f"Here is some information about {topic}.", top_k=100, top_p=0.95, min_length=200)
-	
+    return gpt2_model.generate_one(prompt=f"Here is some information about {topic}.", top_k=50, top_p=0.95, min_length=200)
+
 def generate(topic, wiki=True):
     if wiki:
         try:
-            topic_summary = remove_punctuation(wikipedia.summary(topic))
+            topic_search = wikipedia.search(topic, results=3)
+            topic_summary = remove_punctuation(wikipedia.summary(topic_search[0], auto_suggest=False))
+        except wikipedia.DisambiguationError as e:
+            page = e.options[0]
+            topic_summary = remove_punctuation(wikipedia.summary(page))
         except:
             return(f"Method A struggled to find information about {topic}, please try a different topic!")
     else:
@@ -190,7 +195,7 @@ def generate(topic, wiki=True):
     print(limerick)
 
     return limerick
-	
+    
 def compare_summaries(topic):
     wiki_limerick = generate(topic)
     gpt2_limerick = generate(topic, wiki=False)
@@ -200,9 +205,8 @@ def compare_summaries(topic):
     print(output1 + "\n" + output2)
 
     return output1, output2
-	
-import gradio as gr
-description = "Generates limericks (five-line poems with a rhyme scheme of AABBA) via two different methods. Please be patient as it can take up to a minute to generate both limericks."
+
+description = "Generates limericks (five-line poems with a rhyme scheme of AABBA) via two different methods"
 article = '<center><big><strong>Limerick Generation</strong></big></center>'\
         '<center><strong>By Ans Farooq</strong></center>'\
         '<center><small>Level 4 Individual Project</small></center>'\
@@ -220,14 +224,14 @@ article = '<center><big><strong>Limerick Generation</strong></big></center>'\
         'cohesive limericks automatically, using two different methods. The '\
         'results of this project are to be evaluated through human comparisons.'
 
-input = gr.inputs.Textbox(label='Topic')
-output1 = gr.outputs.Textbox(label='Method A')
-output2 = gr.outputs.Textbox(label='Method B')
+gr_input = gr.inputs.Textbox(label='Topic')
+gr_output1 = gr.outputs.Textbox(label='Method A')
+gr_output2 = gr.outputs.Textbox(label='Method B')
 
 interface = gr.Interface(
     fn=compare_summaries, 
-    inputs=input, 
-    outputs=[output1, output2],
+    inputs=gr_input, 
+    outputs=[gr_output1, gr_output2],
     title="Text-generation with rhyme and rhythm",
     layout="horizontal",
     theme="peach",