Update app.py

app.py

@@ -12,6 +12,7 @@ import fitz
 import PyPDF2
 import gradio
 import sys
+from mistralai import Mistral, DocumentURLChunk, ImageURLChunk, TextChunk, OCRResponse
 from pathlib import Path
 utils_dir = Path(__file__).parent / 'utils'
 sys.path.append(str(utils_dir))
@@ -19,9 +20,13 @@ from openai_utils import *
 import base64
 from pdf2image import convert_from_bytes
 import requests
+import bibtexparser
+from pybtex.database import parse_string
+from pybtex.plugin import find_plugin
+
 PRIVATE_API_KEY = os.getenv('PRIVATE_API_KEY')
 PRIVATE_API_BASE = os.getenv('PRIVATE_API_BASE')
-
+MISTRAL_API = os.getenv('MISTRAL_API')
 
 def insert_sentence(text, sentence, interval):
     lines = text.split('\n')
@@ -44,7 +49,18 @@ def insert_sentence(text, sentence, interval):
         new_lines.append(separator.join(new_words))
 
     return '\n'.join(new_lines)
-    
+
+
+def format_bibtex(paper, style='apa'):
+    bibtex_entry = paper["citationStyles"]["bibtex"]
+    bib_data = parse_string(bibtex_entry, 'bibtex')
+    formatter = find_plugin('pybtex.style.formatting', style)()
+    entries = list(bib_data.entries.values())
+    if not entries:
+        return "No valid entries found."
+    formatted_entry = formatter.format_entries(entries)
+    return '\n'.join(entry.text.render_as('text') for entry in formatted_entry)
+
 def search_paper(query):
     SEMANTIC_SCHOLAR_API_URL = "https://api.semanticscholar.org/graph/v1/paper/"
     url = f"{SEMANTIC_SCHOLAR_API_URL}search?query={query}&limit=3&fields=url,title,abstract&fieldsOfStudy=Computer Science"
@@ -57,10 +73,21 @@ def search_paper(query):
 
     return response.json()    
 
-def split_text_into_chunks(text, chunk_size=300):
-    words = text.split()  
-    chunks = [" ".join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]
-    return chunks
+def get_combined_markdown(pdf_response: OCRResponse) -> str:
+    markdowns: list[str] = []
+    for page in pdf_response.pages:
+        markdowns.append(page.markdown)
+
+    return "\n\n".join(markdowns)
+
+def split_text_into_chunks(pdf_response: OCRResponse) -> str:
+    # words = text.split()  
+    # chunks = [" ".join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]
+    # return chunks
+    markdowns: list[str] = []
+    for page in pdf_response.pages:
+        markdowns.append(page.markdown)
+    return markdowns
 
 def download_pdf(paper):
     pdf_url = paper["openAccessPdf"]["url"]
@@ -70,8 +97,7 @@ def download_pdf(paper):
 
 
         file_object = BytesIO(response.content)
-        extract_text = extract_chapter(file_object)
-        chunks = split_text_into_chunks(extract_text)
+        chunks = extract_chapter(file_object)
         return chunks
     except:
         return []
@@ -79,7 +105,7 @@ def download_pdf(paper):
 
 def recommendation(s2_id, limit=500):
     SEMANTIC_SCHOLAR_API_URL = "https://api.semanticscholar.org/recommendations/v1/papers/forpaper/"
-    url = f"{SEMANTIC_SCHOLAR_API_URL}{s2_id}?limit={limit}&fields=url,title,abstract,publicationDate,isOpenAccess,openAccessPdf"
+    url = f"{SEMANTIC_SCHOLAR_API_URL}{s2_id}?limit={limit}&fields=url,title,abstract,publicationDate,isOpenAccess,openAccessPdf,citationStyles"
 
     # print(url)
     response = requests.get(url)
@@ -92,22 +118,20 @@ def recommendation(s2_id, limit=500):
 
 
 def extract_chapter(file_object):
-    pdf_reader = PyPDF2.PdfReader(file_object)
-    
-
-    num_pages = len(pdf_reader.pages)
-    extraction_started = False
-    extracted_text = ""
-    for page_number in range(num_pages):
-        page = pdf_reader.pages[page_number]
-        page_text = page.extract_text()
-        extraction_started = True
-        page_number_start = page_number
-        if extraction_started:
-            extracted_text += page_text
-            if page_number_start + 1 < page_number:
-                break
-    return extracted_text
+    client = Mistral(api_key=MISTRAL_API)
+    uploaded_file = client.files.upload(
+        file={
+            "file_name": "retrieve.pdf",
+            "content": file_object.read(),
+        },
+        purpose="ocr",
+    )
+
+    signed_url = client.files.get_signed_url(file_id=uploaded_file.id, expiry=1)
+    pdf_response = client.ocr.process(document=DocumentURLChunk(document_url=signed_url.url), model="mistral-ocr-latest", include_image_base64=True)
+    # response_dict = json.loads(pdf_response.json())
+    chunks = split_text_into_chunks(pdf_response)
+    return chunks
         
 
 
@@ -138,7 +162,8 @@ class Reviewer:
                 for paper in papers:
                     retrieval_content += f"Relevant Paper {str(cnt)}:\n"
                     retrieval_content += f"Title: {paper['title']}\n{paper['content']}\n\n"
-                    retrieved_papers += f"{str(cnt)}. {paper['title']}\n"
+                    formatted_citation = format_bibtex(paper, 'unsrt')
+                    retrieved_papers += f"{str(cnt)}. {formatted_citation} ({paper['url']})\n\n"
                     cnt += 1
                 text = retrieval_content + content
                 chat_review_text = self.chat_review(text=text)  
@@ -215,8 +240,8 @@ class Reviewer:
         return rec_papers
 
     def extract_related_content(self, papers, aspect):
-        os.environ["OPENAI_BASE_URL"] = self.api_base
-        os.environ["OPENAI_API_KEY"] = self.api
+        os.environ["OPENAI_BASE_URL"] = PRIVATE_API_BASE
+        os.environ["OPENAI_API_KEY"] = PRIVATE_API_KEY
         client = AsyncOpenAI()
 
         messages = []
@@ -248,7 +273,7 @@ class Reviewer:
             )
         )
 
-        paper_data_list = [{"title": paper["title"], "content": ""} for paper in papers]
+        paper_data_list = [{"title": paper["title"], "content": "", "citationStyles": paper["citationStyles"], "url": paper["url"]} for paper in papers]
 
         for (paper_idx, chunk_idx), response in zip(chunk_index_map, responses):
             if response.strip().lower().startswith("yes"):
@@ -314,7 +339,7 @@ Organize the result in JSON format as follows:
         for paper_data, response in zip(paper_data_list, responses):
             # print(response)
             response = json.loads(response)
-            results.append({"title": paper_data["title"], "content": response["revised_text"]})
+            results.append({"title": paper_data["title"], "content": response["revised_text"], "citationStyles": paper_data["citationStyles"], "url": paper_data["url"]})
         return results
 
 
@@ -372,7 +397,7 @@ Organize the result in JSON format as follows:
                 result = ""
                 limit_cnt = 1
                 for limitation in limitations:
-                    result += f"{str(limit_cnt)}. {limitation}\n"
+                    result += f"{str(limit_cnt)}. {limitation}\n\n"
                     limit_cnt += 1
             # for choice in response.choices:
             #     result += choice.message.content 
@@ -390,7 +415,7 @@ Organize the result in JSON format as follows:
         query = title
         search_results = search_paper(query)
         if search_results != [] and search_results["data"][0]["title"].lower() == title.lower():
-            search_result = search_results[0]
+            search_result = search_results["data"][0]
             retrieval = recommendation(search_result["paperId"])
             recommended_paper_list = []
             for recommended_paper in retrieval["recommendedPapers"]:
@@ -443,7 +468,7 @@ Organize the result in JSON format as follows:
         file_object = BytesIO(pdf_path) # TODO
         pdf_reader = PyPDF2.PdfReader(file_object)
         
-        doc = fitz.open(stream=pdf_path, filetype="pdf") # TODO
+        doc = fitz.open(stream=pdf_path, filetype="pdf") # TODO path/bytes
         page = doc.load_page(0)
         pix = page.get_pixmap()
         image_bytes = pix.tobytes("png")
@@ -470,21 +495,21 @@ Organize the result in JSON format as follows:
         title = response["title"]
         abstract = response["abstract"]
 
+        client = Mistral(api_key=MISTRAL_API)
+        file_object.seek(0) 
+        uploaded_file = client.files.upload(
+            file={
+                "file_name": "upload.pdf",
+                "content": file_object.read(),
+            },
+            purpose="ocr",
+        )
 
+        signed_url = client.files.get_signed_url(file_id=uploaded_file.id, expiry=1)
+        pdf_response = client.ocr.process(document=DocumentURLChunk(document_url=signed_url.url), model="mistral-ocr-latest", include_image_base64=True)
+        # response_dict = json.loads(pdf_response.json())
+        extracted_text = get_combined_markdown(pdf_response)
 
-        num_pages = len(pdf_reader.pages)
-        extraction_started = False
-        extracted_text = ""
-        for page_number in range(num_pages):
-            page = pdf_reader.pages[page_number]
-            page_text = page.extract_text()
-
-            extraction_started = True
-            page_number_start = page_number
-            if extraction_started:
-                extracted_text += page_text
-                if page_number_start + 1 < page_number:
-                    break
         return extracted_text, title, abstract
 
 def main(api,api_base, paper_pdf, aspect, model_name, limit_num, enable_rag):  
@@ -511,9 +536,6 @@ def main(api,api_base, paper_pdf, aspect, model_name, limit_num, enable_rag):
     return retrieved_content, comments, output2
 
 
-
-
-
 ########################################################################################################    
 
 title = "LimitGen"