Spaces:
Running
Running
| 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 | |
| class PaperMetaData: | |
| arxiv_id: str | |
| title: str | |
| summary: str | |
| text: str | |
| def _clean_field(text: str) -> str: | |
| text = re.sub(r"\n", " ", text) | |
| text = re.sub(r"\s+", " ", text) | |
| return text | |
| 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(r"\x03|\x02", "", text) | |
| text = re.sub(r"-\s+", "", text) | |
| text = re.sub(r"\n", " ", text) | |
| return text | |
| 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: | |
| tokenizer = AutoTokenizer.from_pretrained(qa_model_name) | |
| model = AutoModelForQuestionAnswering.from_pretrained(qa_model_name) | |
| qa_pipeline = pipeline("question-answering", model=model, tokenizer=tokenizer) | |
| return qa_pipeline | |
| def get_answer(question: str, context: str) -> str: | |
| qa_pipeline = get_qa_pipeline() | |
| prediction = qa_pipeline(question=question, context=context) | |
| return prediction["answer"] | |
| demo = gr.Blocks() | |
| with demo: | |
| gr.Markdown("# arXiv Paper Q&A\nImport an arXiv paper and ask questions about it!") | |
| gr.Markdown("⚠️ This Space is still Work in progress! ⚠️") | |
| 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" | |
| ) | |
| 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_data, | |
| inputs=arxiv_url, | |
| outputs=[paper_title, paper_summary, paper_text], | |
| ) | |
| gr.Markdown("## 🤨 Ask a question about the paper") | |
| question = gr.Textbox(label="Ask a question about the paper:") | |
| ask_button = gr.Button("Ask me 🤖") | |
| answer = gr.Textbox(label="Answer:") | |
| ask_button.click(fn=get_answer, inputs=[question, paper_summary], outputs=answer) | |
| demo.launch() | |