Create app.py

app.py

@@ -0,0 +1,62 @@
+import torch
+import re
+import PyPDF2
+import utils
+import streamlit as st
+from transformers import BertTokenizerFast, EncoderDecoderModel
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+model_id = 'mrm8488/bert2bert_shared-spanish-finetuned-summarization'
+tokenizer = BertTokenizerFast.from_pretrained(model_id)
+modelo = EncoderDecoderModel.from_pretrained(model_id).to(device)
+
+def generate_summary(text):
+   inputs = tokenizer([text], padding="max_length", truncation=True, max_length=512, return_tensors="pt")
+   input_ids = inputs.input_ids.to(device)
+   attention_mask = inputs.attention_mask.to(device) #attention_mask only says that the model that this words are not pedded
+   output = modelo.generate(input_ids, attention_mask=attention_mask)
+   return tokenizer.decode(output[0], skip_special_tokens=True)
+
+def summarize_pdf(pdf_file):
+    if pdf_file is not None: 
+        with st.spinner('Generando resumen, espera un poco...'):
+            reader = PyPDF2.PdfReader(pdf_file)
+            title = reader.metadata.title
+            author = reader.metadata.author
+            pages =reader.pages
+            text = [pages[i].extract_text() for i in range(len(pages))]
+            text = [utils.drop_non_relevant_text(utils.preprocess_text(x)) for x in text]
+            text = [' '.join(x) for x in text]
+            text=[x+'\n' if len(x) < 50 else generate_summary(x)+'  \n' for x in text]
+            results = [reader.metadata.title+'  \n', reader.metadata.author+'  \n'] + text
+            st.session_state["summary"] = ' '.join(results)
+
+## Graphic interfaz
+def output(pdf_file):
+    if pdf_file is not None: 
+        reader = PyPDF2.PdfReader(pdf_file)
+        title = reader.metadata.title
+        st.session_state["summary"] = title
+
+if 'summary' not in st.session_state:
+    st.session_state['summary'] = ''
+#output = summarize_pdf(pdf_file)
+#reader = PyPDF2.PdfReader(pdf_file)
+#        title = reader.metadata.title
+#        output = title
+#        st.write(output)
+st.caption('Demo para la generación de resumenes en español')
+with st.sidebar:
+    with st.container(border = True):
+        st.title('PDF-Summarizer para español')
+    st.caption('Este demo está basado en el modelo:  \n mrm8488/bert2bert_shared-spanish-finetuned-summarization  \n Creado por Manuel Romero/@mrm8488 con el soporte de Narrativa')
+    pdf_file = st.file_uploader('Carga tu archivo PDF', type="pdf")
+    corre_button = st.button('Genera resumen', 
+                         on_click=summarize_pdf,
+                             args = (pdf_file, ),
+                         help = 'Presiona para generar resumen')
+    #if pdf_file is not None: 
+        
+container = st.container(height=300)
+container.write('Resumen:')
+container.write(st.session_state["summary"])
+