File size: 5,184 Bytes
39d67a2 |
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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
import fitz # PyMuPDF
import json
import os
import re
from sentence_transformers import SentenceTransformer
import pickle
class PDFProcessor:
def __init__(self, pdf_directory="/Users/maraksa/Downloads/chatbot/WebAIM/"):
self.pdf_directory = pdf_directory
self.embedder = SentenceTransformer('all-MiniLM-L6-v2')
# Check if directory exists
if not os.path.exists(pdf_directory):
os.makedirs(pdf_directory)
print(f"Created directory: {pdf_directory}")
print("Please add your WebAIM PDF files to this directory.")
def clean_text(self, text):
"""Clean extracted text from PDF"""
# Remove extra whitespace and line breaks
text = re.sub(r'\s+', ' ', text)
# Remove common PDF artifacts
text = re.sub(r'Page \d+ of \d+', '', text)
text = re.sub(r'WebAIM.*?\n', '', text)
return text.strip()
def extract_text_from_pdf(self, pdf_path):
"""Extract text from PDF with page information"""
print(f"Processing: {os.path.basename(pdf_path)}")
doc = fitz.open(pdf_path)
pages_content = []
for page_num in range(len(doc)):
page = doc[page_num]
text = page.get_text()
# Clean the text
cleaned_text = self.clean_text(text)
# Skip pages with very little content
if len(cleaned_text) < 50:
continue
# Clean and chunk text
chunks = self.chunk_text(cleaned_text, chunk_size=500)
for chunk_idx, chunk in enumerate(chunks):
if len(chunk.strip()) > 30: # Only keep substantial chunks
pages_content.append({
'text': chunk,
'source_file': os.path.basename(pdf_path),
'page_number': page_num + 1,
'chunk_id': chunk_idx,
'source_type': 'WebAIM'
})
doc.close()
print(f"✅ Extracted {len(pages_content)} chunks from {os.path.basename(pdf_path)}")
return pages_content
def chunk_text(self, text, chunk_size=500, overlap=50):
"""Split text into overlapping chunks"""
words = text.split()
chunks = []
for i in range(0, len(words), chunk_size - overlap):
chunk = ' '.join(words[i:i + chunk_size])
if chunk.strip():
chunks.append(chunk.strip())
return chunks
def process_all_pdfs(self):
"""Process all PDFs in the directory"""
all_content = []
# Check if PDFs exist
pdf_files = [f for f in os.listdir(self.pdf_directory) if f.endswith('.pdf')]
if not pdf_files:
print(f"❌ No PDF files found in {self.pdf_directory}")
print("Please add your WebAIM PDF files to the pdfs/ directory")
return []
print(f"Found {len(pdf_files)} PDF files:")
for pdf_file in pdf_files:
print(f" - {pdf_file}")
for filename in pdf_files:
pdf_path = os.path.join(self.pdf_directory, filename)
try:
content = self.extract_text_from_pdf(pdf_path)
all_content.extend(content)
except Exception as e:
print(f"❌ Error processing {filename}: {str(e)}")
return all_content
def create_knowledge_base(self, output_path="knowledge_base.json"):
"""Create searchable knowledge base from PDFs"""
print("🚀 Starting PDF processing...")
all_content = self.process_all_pdfs()
if not all_content:
print("❌ No content extracted. Please check your PDF files.")
return None
print(f"📄 Total chunks extracted: {len(all_content)}")
print("🧠 Creating embeddings... (this may take a few minutes)")
texts = [item['text'] for item in all_content]
embeddings = self.embedder.encode(texts, show_progress_bar=True)
# Save knowledge base
knowledge_base = {
'content': all_content,
'embeddings': embeddings.tolist(),
'metadata': {
'total_chunks': len(all_content),
'embedding_model': 'all-MiniLM-L6-v2',
'chunk_size': 500,
'overlap': 50
}
}
with open(output_path, 'w') as f:
json.dump(knowledge_base, f, indent=2)
print(f"✅ Knowledge base saved to {output_path}")
print(f"📊 Summary:")
print(f" - Total chunks: {len(all_content)}")
print(f" - Embedding dimensions: {len(embeddings[0])}")
print(f" - File size: {os.path.getsize(output_path) / 1024 / 1024:.2f} MB")
return knowledge_base
# Usage
if __name__ == "__main__":
processor = PDFProcessor()
knowledge_base = processor.create_knowledge_base() |