Spaces:
Sleeping
Sleeping
| # classes/pdf_text_extractor.py | |
| import os | |
| import openai | |
| from PyPDF2 import PdfReader | |
| import re | |
| from tqdm import tqdm | |
| from prompts import PDF_SYSTEM_PROMPT | |
| from config import llm_configs | |
| import time | |
| class PDFTextExtractor: | |
| """ | |
| A class to handle PDF text extraction and preprocessing for podcast preparation. | |
| """ | |
| def __init__(self, pdf_path, output_path, model_name="llama3-8b-8192", llm_config=None, max_chars=100000, chunk_size=1000): | |
| """ | |
| Initialize the PDFTextExtractor with paths and model details. | |
| Args: | |
| pdf_path (str): Path to the PDF file. | |
| output_path (str): Path to save the cleaned text file. | |
| model_name (str): Name of the model to use for text processing. | |
| llm_config (dict): Configuration for the LLM. | |
| max_chars (int): Maximum number of characters to process from the PDF. | |
| chunk_size (int): Size of text chunks to process at a time. | |
| """ | |
| self.pdf_path = pdf_path | |
| self.output_path = output_path | |
| self.max_chars = max_chars | |
| self.chunk_size = chunk_size | |
| self.model_name = model_name | |
| self.llm_config = llm_config or llm_configs.get(model_name) | |
| if self.llm_config is None: | |
| raise ValueError(f"Model configuration for {model_name} not found in llm_configs.") | |
| # System prompt for text processing | |
| self.system_prompt = PDF_SYSTEM_PROMPT | |
| def create_client(self): | |
| openai.api_key = self.llm_config["api_key"] | |
| openai.api_base = self.llm_config["base_url"] | |
| return openai | |
| def validate_pdf(self): | |
| """Check if the file exists and is a valid PDF.""" | |
| if not os.path.exists(self.pdf_path): | |
| print(f"Error: File not found at path: {self.pdf_path}") | |
| return False | |
| if not self.pdf_path.lower().endswith('.pdf'): | |
| print("Error: File is not a PDF") | |
| return False | |
| return True | |
| def extract_text(self): | |
| """Extract text from the PDF, limited by max_chars.""" | |
| if not self.validate_pdf(): | |
| return None | |
| with open(self.pdf_path, 'rb') as file: | |
| pdf_reader = PdfReader(file) | |
| num_pages = len(pdf_reader.pages) | |
| print(f"Processing PDF with {num_pages} pages...") | |
| extracted_text = [] | |
| total_chars = 0 | |
| for page_num in range(num_pages): | |
| page = pdf_reader.pages[page_num] | |
| text = page.extract_text() or "" | |
| if total_chars + len(text) > self.max_chars: | |
| remaining_chars = self.max_chars - total_chars | |
| extracted_text.append(text[:remaining_chars]) | |
| print(f"Reached {self.max_chars} character limit at page {page_num + 1}") | |
| break | |
| extracted_text.append(text) | |
| total_chars += len(text) | |
| print(f"Processed page {page_num + 1}/{num_pages}") | |
| final_text = '\n'.join(extracted_text) | |
| print(f"Extraction complete! Total characters: {len(final_text)}") | |
| return final_text | |
| def create_word_bounded_chunks(self, text): | |
| """Split text into chunks around the target size.""" | |
| words = text.split() | |
| chunks = [] | |
| current_chunk = [] | |
| current_length = 0 | |
| for word in words: | |
| word_length = len(word) + 1 # +1 for the space | |
| if current_length + word_length > self.chunk_size and current_chunk: | |
| chunks.append(' '.join(current_chunk)) | |
| current_chunk = [word] | |
| current_length = word_length | |
| else: | |
| current_chunk.append(word) | |
| current_length += word_length | |
| if current_chunk: | |
| chunks.append(' '.join(current_chunk)) | |
| return chunks | |
| def process_chunk(self, text_chunk): | |
| """Process a text chunk with the model and return the cleaned text.""" | |
| conversation = [ | |
| {"role": "system", "content": self.system_prompt}, | |
| {"role": "user", "content": text_chunk} | |
| ] | |
| client = self.create_client() | |
| response = client.ChatCompletion.create( | |
| model=self.model_name, | |
| messages=conversation, | |
| ) | |
| processed_text = response.choices[0].message.content | |
| return processed_text | |
| def clean_and_save_text(self): | |
| """Extract, clean, and save processed text to a file.""" | |
| extracted_text = self.extract_text() | |
| if not extracted_text: | |
| return None | |
| chunks = self.create_word_bounded_chunks(extracted_text) | |
| processed_text = "" | |
| with open(self.output_path, 'w', encoding='utf-8') as out_file: | |
| for chunk_num, chunk in enumerate(tqdm(chunks, desc="Processing chunks")): | |
| processed_chunk = self.process_chunk(chunk) | |
| processed_text += processed_chunk + "\n" | |
| out_file.write(processed_chunk + "\n") | |
| out_file.flush() | |
| time.sleep(3) # To avoid rate limiting | |
| print(f"\nExtracted and cleaned text has been saved to {self.output_path}") | |
| return self.output_path | |