Spaces:
Sleeping
Sleeping
File size: 2,906 Bytes
9aa78e6 50cbe29 9aa78e6 d22c38f 9aa78e6 1c322e9 9aa78e6 d22c38f 9aa78e6 1c322e9 9aa78e6 ad329be 9aa78e6 d22c38f 9aa78e6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
from ..utils.functions import getConfig, convertImageToBase64, getYaml
from pdf2image import convert_from_bytes
from ..utils.logger import logger
from dotenv import load_dotenv
from openai import OpenAI
from PIL import Image
import math
import os
load_dotenv()
class ExtractPdfDetails:
def __init__(self):
logger.info("INITIALIZING EXTRACT PDF DETAILS")
self.config = getConfig(os.path.join(os.getcwd(), "config.ini"))
self.prompts = getYaml(os.path.join(os.getcwd(), "prompts.yaml"))
self.llmClient = OpenAI(
base_url = self.config.get("GROQ CONFIG", "BASEURL"),
api_key = os.environ["GROQ_API_KEY"]
)
def convertToImages(self, pdfBytes: str) -> list[Image.Image]:
"""
Convert a pdf to a list of images
Args:
pdfBytes: bytes of the pdf file
Returns:
list[Image.Image]: list of pdf pages as images
"""
try:
logger.info(f"Converting pdf to images")
images = convert_from_bytes(pdfBytes)
return images
except Exception as e:
logger.exception(f"Error converting pdf to images: {e}")
return None
def chunkImages(self, images: list[Image.Image]) -> list[list[Image.Image]]:
"""
Chunk the images into smaller chunks
Args:
images: list of images
Returns:
chunks: list of chunks of images
"""
try:
logger.info("Chunking the images")
batchSize = self.config.getint("DETAIL EXTRACTOR", "BATCHSIZE")
nBatches = math.ceil(len(images) / batchSize)
chunks = [images[batchSize * x: batchSize * x + batchSize] for x in range(nBatches)]
return chunks
except Exception as e:
logger.exception(f"Error chunking the images: {e}")
return None
def extractDetailsFromChunk(self, images: list[Image.Image]) -> str:
"""
Extract details from a chunk of images
Args:
images: list of images
Returns:
details: string of details extracted from the images
"""
try:
logger.info("Extracting details from the images")
completion = self.llmClient.chat.completions.create(
model = self.config.get("DETAIL EXTRACTOR", "VLM"),
messages = [
{"role": "system", "content": self.prompts["detailExtractorPrompt"]},
{"role": "user", "content": [{"type": "image_url", "image_url": {"url": convertImageToBase64(image)}} for image in images]}
]
)
response = completion.choices[0].message.content
return response
except Exception as e:
logger.exception(f"Error extracting details from the images: {e}")
return None |