Update app.py

app.py

@@ -1,57 +1,13 @@
 import streamlit as st
 import io
 import requests
-import pdfplumber
-import os
-from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
 
-from huggingface_hub import login
+from docquery import document, pipeline
 
-token = os.getenv("hf_token")
-login(token=token)
 
-batch_size = 96
-n_epochs = 2
-base_LM_model = "roberta-base"
-max_seq_len = 512
-learning_rate = 3e-5
-warmup_proportion = 0.2
-doc_stride=128
-max_query_length=64
 
 
-
-
-
-def AImodel(text,questionText):
-    model_name = "deepset/roberta-base-squad2"
-    
-
-    # a) Get predictions
-    nlp = pipeline('question-answering', model=model_name, tokenizer=model_name)
-    QA_input = {
-        'question': questionText,
-        'context': text
-    }
-    res = nlp(QA_input)
-
-    
-
-    # b) Load model & tokenizer
-    model = AutoModelForQuestionAnswering.from_pretrained(model_name)
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    inputs = tokenizer(
-        questionText, text,
-        add_special_tokens=True,
-        max_length=max_seq_len,
-        truncation=True,
-        return_tensors="pt",
-        padding="max_length",
-        stride=doc_stride,
-        return_overflowing_tokens=True,
-        return_offsets_mapping=True
-    )
-    return(res)
+p = pipeline('document-question-answering')
 
 
 
@@ -63,36 +19,12 @@ def fextractURL(pdf_path):
             # If the URL ends with .pdf, use pdfplumber directly
             r = requests.get(pdf_path)
             f = io.BytesIO(r.content)
-            with pdfplumber.open(f) as pdf:
-                for page in pdf.pages:
-                    extracted_data += page.extract_text() + "\n"  # Extract text
-                    tables = page.extract_tables()  # Extract tables
-                    for table in tables:
-                        for row in table:
-                            extracted_data += "\t".join(str(cell) for cell in row) + "\n"
-        else:
-            # If the URL does not end with .pdf, download the PDF first
-            response = requests.get(pdf_path)
-            pdf_content = response.content
-
-            # Save the PDF locally
-            pdf_filename = 'downloaded_document.pdf'
-            with open(pdf_filename, 'wb') as pdf_file:
-                pdf_file.write(pdf_content)
-
-            # Extract content using pdfplumber
-            with pdfplumber.open(pdf_filename) as pdf:
-                for page in pdf.pages:
-                    extracted_data += page.extract_text() + "\n"  # Extract text
-                    tables = page.extract_tables()  # Extract tables
-                    for table in tables:
-                        for row in table:
-                            extracted_data += "\t".join(str(cell) for cell in row) + "\n"
-
-            # Delete the PDF file
-            os.remove(pdf_filename)
+         
+            doc = document.load_document(f)
+            for q in ["What is the 2022 net income?", "What is the 2023 net income ?"]:
+                extracted_data+= (q, p(question=q, **doc.context))
     except Exception as e:
-        st.error(f"An error occurred: {str(e)}")
+        st.error(f"An error o0000ccurred: {str(e)}")
 
     return extracted_data
 
@@ -102,19 +34,13 @@ st.markdown(vert_space, unsafe_allow_html=True)
 st.write("Extract full text from PDF URL")
 
 pdfURL = st.text_input(label="PDF URL", value="", max_chars=None, key=None, type="default", help=None, autocomplete=None, on_change=None, args=None, kwargs=None, placeholder=None, disabled=False, label_visibility="visible")
-questionText = st.text_input(label="question", value="", max_chars=None, key=None, type="default", help=None, autocomplete=None, on_change=None, args=None, kwargs=None, placeholder=None, disabled=False, label_visibility="visible")
 button = st.button(label='Extract', key=None, help=None, on_click=None, args=None, kwargs=None, type="secondary", disabled=False, use_container_width=False)
 extractedText = st.empty()
 
 if button:
     try:
         text = fextractURL(pdfURL)
-        AItext = AImodel(text,questionText)
-        extractedText.text(AItext)
-
+        print(text)
+        extractedText.text(text)
     except Exception as e:
-        st.error(f"An error occurred: {str(e)}")
-
-
-
-
+        st.error(f"An error occurrrrred: {str(e)}")