Tuana/PDF-Summarizer

1import streamlit as st2from haystack.document_stores import InMemoryDocumentStore3from haystack.nodes import TransformersSummarizer, PreProcessor, PDFToTextConverter, Crawler4from haystack.schema import Document5import logging6import base647from PIL import Image8import validators9 10@st.cache(hash_funcs={"builtins.SwigPyObject": lambda _: None},allow_output_mutation=True)11def start_haystack():12    document_store = InMemoryDocumentStore()13    preprocessor = PreProcessor(14        clean_empty_lines=True,15        clean_whitespace=True,16        clean_header_footer=True,17        split_by="word",18        split_length=200,19        split_respect_sentence_boundary=True,20    )21    summarizer = TransformersSummarizer(model_name_or_path="facebook/bart-large-cnn")22    return document_store, summarizer, preprocessor23 24 25def pdf_to_document_store(pdf_file):26    document_store.delete_documents()27    converter = PDFToTextConverter(remove_numeric_tables=True, valid_languages=["en"])28    with open("temp-path.pdf", 'wb') as temp_file:29        base64_pdf = base64.b64encode(pdf_file.read()).decode('utf-8')30        temp_file.write(base64.b64decode(base64_pdf))31        doc = converter.convert(file_path="temp-path.pdf", meta=None)32        preprocessed_docs=preprocessor.process(doc)33        document_store.write_documents(preprocessed_docs)34        temp_file.close()35    36def summarize(content):37    pdf_to_document_store(content)38    summaries = summarizer.predict(documents=document_store.get_all_documents(), generate_single_summary=True)39    return summaries40 41def set_state_if_absent(key, value):42    if key not in st.session_state:43        st.session_state[key] = value44        45set_state_if_absent("summaries", None)46        47document_store, summarizer, preprocessor = start_haystack()48 49st.title('TL;DR with Haystack')50image = Image.open('header-image.png')51st.image(image)52 53st.markdown( """54This Summarization demo uses a [Haystack TransformerSummarizer node](https://haystack.deepset.ai/pipeline_nodes/summarizer). You can upload a PDF file, which will be converted to text with the [Haystack PDFtoTextConverter](https://haystack.deepset.ai/reference/file-converters#pdftotextconverter). In this demo, we produce 1 summary for the whole file you upload. So, the TransformerSummarizer treats the whole thing as one string, which means along with the model limitations, PDFs that have a lot of unneeded text at the beginning produce poor results. For best results, upload a document that has minimal intro and tables at the top. 55""", unsafe_allow_html=True)56 57uploaded_file = st.file_uploader("Choose a PDF file", accept_multiple_files=False)58                59if uploaded_file is not None :60    if st.button('Summarize Document'):61        with st.spinner("📚    Please wait while we produce a summary..."):62            try:63                st.session_state.summaries = summarize(uploaded_file)64            except Exception as e:65                logging.exception(e)66 67if st.session_state.summaries:68    st.write('## Summary')69    for count, summary in enumerate(st.session_state.summaries):70        st.write(summary.content)71