Update app.py

app.py

@@ -24,7 +24,7 @@ with open("wordFrequency.txt", 'r') as f:
         line = f.readline()
 
 def filter_rhymes(word):
-    filter_list = ['an', 'to', 'on', 'has', 'but', 'the', 'in', 'and', 'a', 'are', 'or', 'its', 'it''s']
+    filter_list = ['an', 'to', 'on', 'has', 'but', 'the', 'in', 'and', 'a', 'are', 'or', 'its', 'it''s'] 
     if word in filter_list:
         return False
     else:
@@ -33,7 +33,7 @@ def filter_rhymes(word):
 def remove_punctuation(text):
     text = re.sub(r'[^\w\s]', '', text)
     text = text.replace("\n", " ")
-    return text
+    return text.strip()
 
 def get_rhymes(inp, level):
     entries = nltk.corpus.cmudict.entries()
@@ -53,7 +53,6 @@ def get_inputs_length(input):
     return len(input_ids)
     
 def get_prediction(sent):
-    
     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 ]
@@ -90,7 +89,7 @@ def get_line(prompt, inputs_len):
         clean_up_tokenization_spaces=True,
         return_full_text=False
     )
-    return remove_punctuation(output[0]['generated_text']).strip()
+    return remove_punctuation(output[0]['generated_text'])
 
 def get_rhyming_line(prompt, rhyming_word, inputs_len):
     output = gpt2_pipeline(
@@ -100,7 +99,7 @@ def get_rhyming_line(prompt, rhyming_word, inputs_len):
         clean_up_tokenization_spaces=True,
         return_full_text=False
     )
-    gpt2_sentence = remove_punctuation(output[0]['generated_text']).strip()
+    gpt2_sentence = remove_punctuation(output[0]['generated_text'])
     while len(gpt2_sentence) == 0:
         output = gpt2_pipeline(
             prompt + ".",
@@ -109,7 +108,7 @@ def get_rhyming_line(prompt, rhyming_word, inputs_len):
             clean_up_tokenization_spaces=True,
             return_full_text=False
         )
-        gpt2_sentence = remove_punctuation(output[0]['generated_text']).strip()
+        gpt2_sentence = remove_punctuation(output[0]['generated_text'])
     
     print(f"\nGetting rhyming line starting with '{gpt2_sentence}' and ending with rhyming word '{rhyming_word}'")
     sentence = gpt2_sentence + " ___ ___ ___ " + rhyming_word
@@ -130,12 +129,12 @@ def get_rhyming_line(prompt, rhyming_word, inputs_len):
 def gpt2_summary(topic):
     output = gpt2_pipeline(
         f"Here is some information about {topic}.",
-        min_length=200,
+        min_length=100,
         max_length=300,
         clean_up_tokenization_spaces=True,
         return_full_text=False
     )
-    return remove_punctuation(output[0]['generated_text']).strip()
+    return remove_punctuation(output[0]['generated_text'])
     
 def generate(topic, wiki=True):
     if wiki:
@@ -144,7 +143,6 @@ def generate(topic, wiki=True):
             print(f"Wikipedia search results for {topic} are: {topic_search}")
             topic_summary = remove_punctuation(wikipedia.summary(topic_search[0], auto_suggest=False))
         except wikipedia.DisambiguationError as e:
-            print("===================== DISAMBIGUATION ERROR =====================")
             print(f"Wikipedia returned a disambiguation error for {topic}. Selecting the first option {e.options[0]} instead.")
             page = e.options[0]
             topic_summary = remove_punctuation(wikipedia.summary(page, auto_suggest=False))