added feature to parse all text from paper pdf

app.py

@@ -1,24 +1,84 @@
+from __future__ import annotations
+
 import re
+from dataclasses import dataclass
+from typing import Tuple
 
 import gradio as gr
 import requests
 import xmltodict
+from PyPDF2 import PdfReader
 from transformers import AutoModelForQuestionAnswering, AutoTokenizer, pipeline
 from transformers.pipelines.question_answering import QuestionAnsweringPipeline
 
 QA_MODEL_NAME = "ixa-ehu/SciBERT-SQuAD-QuAC"
+TEMP_PDF_PATH = "/tmp/arxiv_paper.pdf"
+ARXIV_URL_PATTERN = r"(http|https)://(arxiv.org/pdf/)+([0-9]+\.[0-9]+)\.pdf"
+
+
+def is_valid_url(url: str) -> bool:
+    return re.fullmatch(ARXIV_URL_PATTERN, url) is not None
+
+
+@dataclass
+class PaperMetaData:
+    arxiv_id: str
+    title: str
+    summary: str
+    text: str
+
+    @staticmethod
+    def _clean_field(text: str) -> str:
+        text = re.sub(r"\n", " ", text)
+        text = re.sub(r"\s+", " ", text)
+        return text
+
+    @classmethod
+    def from_api(cls, arxiv_id: str, text: str) -> PaperMetaData:
+        paper_url = f"http://export.arxiv.org/api/query?id_list={arxiv_id}"
+        response = requests.get(paper_url)
+        paper_dict = xmltodict.parse(response.content)["feed"]["entry"]
+        return PaperMetaData(
+            arxiv_id=arxiv_id,
+            title=cls._clean_field(paper_dict["title"]),
+            summary=cls._clean_field(paper_dict["summary"]),
+            text=text,
+        )
 
 
 def clean_text(text: str) -> str:
-    text = re.sub("\n", " ", text)
+    text = re.sub(r"\x03|\x02", "", text)
+    text = re.sub(r"-\s+", "", text)
+    text = re.sub(r"\n", " ", text)
     return text
 
 
-def get_paper_summary(arxiv_id: str) -> str:
-    paper_url = f"http://export.arxiv.org/api/query?id_list={arxiv_id}"
-    response = requests.get(paper_url)
-    paper_dict = xmltodict.parse(response.content)["feed"]["entry"]
-    return clean_text(paper_dict["summary"])
+class PDFPaper:
+    def __init__(self, url: str):
+        if not is_valid_url(url):
+            raise ValueError("The URL provided is not a valid arxiv PDF url.")
+        self.url = url
+        self.arxiv_id = re.fullmatch(ARXIV_URL_PATTERN, url).group(3)
+
+    def _download(self, download_path: str = TEMP_PDF_PATH) -> None:
+        pdf_r = requests.get(self.url)
+        pdf_r.raise_for_status()
+        with open(download_path, "wb") as pdf_file:
+            pdf_file.write(pdf_r.content)
+
+    def read_text(self, pdf_path: str = TEMP_PDF_PATH) -> str:
+        self._download(pdf_path)
+        reader = PdfReader(pdf_path)
+        pdf_text = " ".join([page.extract_text() for page in reader.pages])
+        return clean_text(pdf_text)
+
+    def get_paper_full_data(self) -> PaperMetaData:
+        return PaperMetaData.from_api(arxiv_id=self.arxiv_id, text=self.read_text())
+
+
+def get_paper_data(url: str) -> Tuple[str, str, str]:
+    paper_data = PDFPaper(url=url).get_paper_full_data()
+    return paper_data.title, paper_data.summary, paper_data.text
 
 
 def get_qa_pipeline(qa_model_name: str = QA_MODEL_NAME) -> QuestionAnsweringPipeline:
@@ -36,24 +96,27 @@ def get_answer(question: str, context: str) -> str:
 
 demo = gr.Blocks()
 
-
 with demo:
-    gr.Markdown("# Document QA")
+    gr.Markdown("# arXiv Paper Q&A\nImport an arXiv paper and ask questions about it!")
 
-    # Retrieve paper
-    arxiv_id = gr.Textbox(
-        label="arXiv Paper ID", placeholder="Insert here the ID of a paper on arXiv"
+    gr.Markdown("## 📄 Import the paper on arXiv")
+    arxiv_url = gr.Textbox(
+        label="arXiv Paper URL", placeholder="Insert here the URL of a paper on arXiv"
     )
-    paper_summary = gr.Textbox(label="Paper summary")
-    fetch_document_button = gr.Button("Get Summary")
+    fetch_document_button = gr.Button("Import Paper")
+    paper_title = gr.Textbox(label="Paper Title")
+    paper_summary = gr.Textbox(label="Paper Summary")
+    paper_text = gr.Textbox(label="Paper Text")
     fetch_document_button.click(
-        fn=get_paper_summary, inputs=arxiv_id, outputs=paper_summary
+        fn=get_paper_data,
+        inputs=arxiv_url,
+        outputs=[paper_title, paper_summary, paper_text],
     )
 
-    # QA on paper
+    gr.Markdown("## 🤨 Ask a question about the paper")
     question = gr.Textbox(label="Ask a question about the paper:")
-    answer = gr.Textbox("Answer:")
     ask_button = gr.Button("Ask me 🤖")
+    answer = gr.Textbox(label="Answer:")
     ask_button.click(fn=get_answer, inputs=[question, paper_summary], outputs=answer)