Spaces:
Build error
Build error
| import base64 | |
| import os | |
| import requests | |
| from io import BytesIO | |
| from openai import OpenAI | |
| from pdf2image import convert_from_path | |
| from langchain.schema import Document | |
| from modules.config.constants import TIMEOUT | |
| class GPTParser: | |
| """ | |
| This class uses OpenAI's GPT-4o mini model to parse PDFs and extract text, images and equations. | |
| It is the most advanced parser in the system and is able to handle complex formats and layouts | |
| """ | |
| def __init__(self): | |
| self.client = OpenAI() | |
| self.api_key = os.getenv("OPENAI_API_KEY") | |
| self.prompt = """ | |
| The provided documents are images of PDFs of lecture slides of deep learning material. | |
| They contain LaTeX equations, images, and text. | |
| The goal is to extract the text, images and equations from the slides and convert everything to markdown format. Some of the equations may be complicated. | |
| The markdown should be clean and easy to read, and any math equation should be converted to LaTeX, between $$. | |
| For images, give a description and if you can, a source. Separate each page with '---'. | |
| Just respond with the markdown. Do not include page numbers or any other metadata. Do not try to provide titles. Strictly the content. | |
| """ | |
| def parse(self, pdf_path): | |
| images = convert_from_path(pdf_path) | |
| encoded_images = [self.encode_image(image) for image in images] | |
| chunks = [encoded_images[i : i + 5] for i in range(0, len(encoded_images), 5)] | |
| headers = { | |
| "Content-Type": "application/json", | |
| "Authorization": f"Bearer {self.api_key}", | |
| } | |
| output = "" | |
| for chunk_num, chunk in enumerate(chunks): | |
| content = [ | |
| { | |
| "type": "image_url", | |
| "image_url": {"url": f"data:image/jpeg;base64,{image}"}, | |
| } | |
| for image in chunk | |
| ] | |
| content.insert(0, {"type": "text", "text": self.prompt}) | |
| payload = { | |
| "model": "gpt-4o-mini", | |
| "messages": [{"role": "user", "content": content}], | |
| } | |
| response = requests.post( | |
| "https://api.openai.com/v1/chat/completions", | |
| headers=headers, | |
| json=payload, | |
| timeout=TIMEOUT, | |
| ) | |
| resp = response.json() | |
| chunk_output = ( | |
| resp["choices"][0]["message"]["content"] | |
| .replace("```", "") | |
| .replace("markdown", "") | |
| .replace("````", "") | |
| ) | |
| output += chunk_output + "\n---\n" | |
| output = output.split("\n---\n") | |
| output = [doc for doc in output if doc.strip() != ""] | |
| documents = [ | |
| Document(page_content=page, metadata={"source": pdf_path, "page": i}) | |
| for i, page in enumerate(output) | |
| ] | |
| return documents | |
| def encode_image(self, image): | |
| buffered = BytesIO() | |
| image.save(buffered, format="JPEG") | |
| return base64.b64encode(buffered.getvalue()).decode("utf-8") | |