Update app.py

app.py

@@ -2,6 +2,7 @@ from transformers import T5ForConditionalGeneration,T5Tokenizer
 from transformers import AutoModelWithLMHead, AutoTokenizer
 from transformers import pipeline
 import streamlit as st 
+import re
 
 model = T5ForConditionalGeneration.from_pretrained("Michau/t5-base-en-generate-headline")
 tokenizer = T5Tokenizer.from_pretrained("Michau/t5-base-en-generate-headline")
@@ -9,37 +10,55 @@ tokenizer = T5Tokenizer.from_pretrained("Michau/t5-base-en-generate-headline")
 mrm_tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-summarize-news")
 mrm_model = AutoModelWithLMHead.from_pretrained("mrm8488/t5-base-finetuned-summarize-news")
 
+jules_tokenizer = AutoTokenizer.from_pretrained("JulesBelveze/t5-small-headline-generator")
+jules_model = T5ForConditionalGeneration.from_pretrained("JulesBelveze/t5-small-headline-generator")
+# rouge = Rouge()
+WHITESPACE_HANDLER = lambda k: re.sub('\s+', ' ', re.sub('\n+', ' ', k.strip()))
 
 def generate_title(article):
     text =  "headline: " + article
     encoding = tokenizer.encode_plus(text, return_tensors = "pt", max_length=2048, truncation=True)
     input_ids = encoding["input_ids"]
     attention_masks = encoding["attention_mask"]
-
     beam_outputs = model.generate(
-        input_ids = input_ids,
-        attention_mask = attention_masks,
-        max_length = 50,
-        num_beams = 3,
-        do_sample = True,
-        top_k=10,
-        early_stopping = False,
-    )
+            input_ids = input_ids,
+            attention_mask = attention_masks,
+            max_length = 50,
+            num_beams = 3,
+            do_sample = False,
+            # top_k=10,
+            early_stopping = False,
+        )
 
     return tokenizer.decode(beam_outputs[0])
 
-# def generate_summary(article):
-#   input_ids = mrm_tokenizer.encode(article, return_tensors="pt", add_special_tokens=True)
-
-#   generated_ids = mrm_model.generate(input_ids=input_ids, num_beams=3, max_length=200,  repetition_penalty=2.5, length_penalty=1.0, early_stopping=False, truncation=True)
+def generate_title_2(article):
+    input_ids = tokenizer(
+    [WHITESPACE_HANDLER(article)],
+    return_tensors="pt",
+    padding="max_length",
+    truncation=True,
+    max_length=384
+    )["input_ids"]
 
-#   preds = [mrm_tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True) for g in generated_ids]
+    output_ids = model.generate(
+        input_ids=input_ids,
+        max_length=84,
+        no_repeat_ngram_size=2,
+        num_beams=4
+    )[0]
+    summary = tokenizer.decode(
+    output_ids,
+    skip_special_tokens=True,
+    clean_up_tokenization_spaces=False
+    )
+    return summary
 
-#   return preds[0]
 def generate_summary(article):
     article = article[:1024]
     summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
     return summarizer(article, max_length=130, min_length=30, do_sample=False)
+    
 def main():
     st.title("Text Summarization")
     text = st.text_area("Enter your text here:", "")
@@ -49,11 +68,15 @@ def main():
             st.error("Please enter some text.")
         else:
             title = generate_title(text)
+            title_2 = generate_title_2(text)
             summary = generate_summary(text)
             # summary = summary[0]['summary_text']
 
             st.subheader("Generated Title:")
             st.write(title.replace('<pad>', '').replace('</s>', ''))
+
+            st.subheader("Second Title:")
+            st.write(title_2)
             
             st.subheader("Generated Description:")