Spaces:
Sleeping
Sleeping
| import os | |
| import io | |
| from pdfminer.high_level import extract_text_to_fp | |
| from pdfminer.layout import LAParams | |
| import re | |
| from tqdm import tqdm | |
| import time | |
| from nylon import ChatDatabase, get_keywords | |
| def extract_text_from_pdf(pdf_path): | |
| output_string = io.StringIO() | |
| with open(pdf_path, 'rb') as fin: | |
| laparams = LAParams(line_margin=1.5, char_margin=2.0, word_margin=0.1) | |
| extract_text_to_fp(fin, output_string, laparams=laparams, | |
| output_type='text', codec='utf-8') | |
| return output_string.getvalue() | |
| def process_text_into_paragraphs(text): | |
| # Remove page numbers and headers/footers | |
| text = re.sub(r'\n\d+\n', '\n', text) | |
| text = re.sub(r'^\s*\d+\s*$', '', text, flags=re.MULTILINE) | |
| # Split text into paragraphs | |
| paragraphs = list(set(re.split(r'\n{2,}', text))) | |
| # Clean up each paragraph | |
| cleaned_paragraphs = [] | |
| for para in paragraphs: | |
| # Remove extra whitespace and join broken words | |
| cleaned_para = re.sub(r'\s+', ' ', para).strip() | |
| cleaned_para = re.sub(r'(\w+)-\s*(\w+)', r'\1\2', cleaned_para) | |
| cleaned_para = re.sub(r'(\w+)\s*\n\s*(\w+)', r'\1 \2', cleaned_para) | |
| if cleaned_para: # Only add non-empty paragraphs | |
| cleaned_paragraphs.append(cleaned_para) | |
| return cleaned_paragraphs | |
| def process_pdfs(directory, db): | |
| fixed_timestamp = "2024-10-22 12:00:00" | |
| sender = "Arcana" # Set sender to "Arcana" for all messages | |
| pdf_files = [f for f in os.listdir(directory) if f.endswith('.pdf')] | |
| total_files = len(pdf_files) | |
| with tqdm(total=total_files, desc="Processing PDFs", unit="file") as pbar: | |
| for filename in pdf_files: | |
| pdf_path = os.path.join(directory, filename) | |
| tag = os.path.splitext(filename)[0] # Use filename without .pdf as tag | |
| text = extract_text_from_pdf(pdf_path) | |
| paragraphs = process_text_into_paragraphs(text) | |
| for paragraph in paragraphs: | |
| #print(paragraph) | |
| db.add_message(sender, fixed_timestamp, str(paragraph), tag) | |
| pbar.update(1) | |
| pbar.set_postfix({"Current File": filename}) | |
| print(filename) | |
| def main(foldername): | |
| db_filename = foldername | |
| if os.path.exists(db_filename): | |
| db_filename += '.txt' | |
| print(f"Database file '{db_filename}' already exists. Rewriting existing database...") | |
| #db = ChatDatabase(db_filename) | |
| #else: | |
| print(f"Creating new database '{db_filename}'...") | |
| db = ChatDatabase(db_filename) | |
| pdf_directory = foldername | |
| start_time = time.time() | |
| process_pdfs(pdf_directory, db) | |
| end_time = time.time() | |
| total_time = end_time - start_time | |
| print(f"\nDatabase creation complete. Total time: {total_time:.2f} seconds") | |
| if __name__ == "__main__": | |
| main() | |