Spaces:
Running
Running
File size: 5,738 Bytes
641b53a 10ea2c4 641b53a 10ea2c4 641b53a ce71763 641b53a 10ea2c4 641b53a ce71763 641b53a ce71763 641b53a ce71763 641b53a ce71763 641b53a ce71763 641b53a ce71763 641b53a 0315b16 641b53a | 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 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | '''
Multi-format support (PDF, DOCX, TXT, MD, HTML)
- Intelligent chunking with overlap
- Metadata extraction (title, author, date, file type)
- Text cleaning and normalization
- Duplicate detection
'''
import hashlib
import os
import re
from datetime import datetime
from typing import Dict, List, Optional
import PyPDF2
import tiktoken
from bs4 import BeautifulSoup
from docx import Document
class _RegexTokenizer:
"""Offline fallback tokenizer when tiktoken encoding cannot be loaded."""
_token_pattern = re.compile(r"\w+|[^\w\s]", re.UNICODE)
def encode(self, text: str) -> List[str]:
return self._token_pattern.findall(text)
def decode(self, token_ids: List[str]) -> str:
if not token_ids:
return ""
return " ".join(token_ids)
class DocumentProcessor:
def __init__(self, chunk_size: int = 600, overlap: int = 100, tokenizer_name: str = "gpt2"):
if chunk_size <= 0:
raise ValueError("chunk_size must be > 0")
if overlap < 0:
raise ValueError("overlap must be >= 0")
if overlap >= chunk_size:
raise ValueError("overlap must be smaller than chunk_size")
self.chunk_size = chunk_size
self.overlap = overlap
self.tokenizer_name = tokenizer_name
try:
self._tokenizer = tiktoken.get_encoding(tokenizer_name)
except Exception:
self._tokenizer = _RegexTokenizer()
self._seen_hashes: set = set()
def process_document(self, file_path: str) -> Optional[Dict]:
text = self.extract_text(file_path)
if self._is_duplicate(text):
return None
metadata = self.extract_metadata(file_path)
cleaned_text = self.clean_text(text)
chunks = self.chunk_text(cleaned_text)
return {
'metadata': metadata,
'chunks': chunks,
}
def _is_duplicate(self, text: str) -> bool:
digest = hashlib.sha256(text.encode('utf-8')).hexdigest()
if digest in self._seen_hashes:
return True
self._seen_hashes.add(digest)
return False
def extract_text(self, file_path: str) -> str:
ext = os.path.splitext(file_path)[1].lower()
extractors = {
'.pdf': self._extract_pdf_text,
'.docx': self._extract_docx_text,
'.txt': self._extract_plain_text,
'.md': self._extract_plain_text,
'.html': self._extract_html_text,
}
extractor = extractors.get(ext)
if extractor is None:
raise ValueError(f"Unsupported file type: {ext!r}")
return extractor(file_path)
def extract_metadata(self, file_path: str) -> Dict:
ext = os.path.splitext(file_path)[1].lower()
base = {
'title': os.path.basename(file_path),
'author': 'Unknown',
'date': None,
'file_type': ext,
}
if ext == '.pdf':
base.update(self._pdf_metadata(file_path))
elif ext == '.docx':
base.update(self._docx_metadata(file_path))
if base['date'] is None:
base['date'] = datetime.now().isoformat()
return base
def clean_text(self, text: str) -> str:
text = re.sub(r'\s+', ' ', text)
return text.strip()
def chunk_text(self, text: str) -> List[str]:
token_ids = self._tokenizer.encode(text)
if not token_ids:
return []
chunks: List[str] = []
step = self.chunk_size - self.overlap
start = 0
while start < len(token_ids):
end = min(start + self.chunk_size, len(token_ids))
chunk = self._tokenizer.decode(token_ids[start:end])
if chunk.strip():
chunks.append(chunk)
start += step
return chunks
def count_tokens(self, text: str) -> int:
return len(self._tokenizer.encode(text))
# --- private extractors ---
def _extract_pdf_text(self, file_path: str) -> str:
with open(file_path, 'rb') as f:
reader = PyPDF2.PdfReader(f)
return ''.join(page.extract_text() or '' for page in reader.pages)
def _extract_docx_text(self, file_path: str) -> str:
doc = Document(file_path)
return '\n'.join(para.text for para in doc.paragraphs)
def _extract_plain_text(self, file_path: str) -> str:
with open(file_path, 'r', encoding='utf-8') as f:
return f.read()
def _extract_html_text(self, file_path: str) -> str:
with open(file_path, 'r', encoding='utf-8') as f:
soup = BeautifulSoup(f, 'html.parser')
return soup.get_text(separator=' ')
# --- private metadata helpers ---
def _pdf_metadata(self, file_path: str) -> Dict:
result = {}
try:
with open(file_path, 'rb') as f:
info: Dict = dict(PyPDF2.PdfReader(f).metadata or {})
if info.get('/Title'):
result['title'] = info['/Title']
if info.get('/Author'):
result['author'] = info['/Author']
if info.get('/CreationDate'):
result['date'] = info['/CreationDate']
except Exception:
pass
return result
def _docx_metadata(self, file_path: str) -> Dict:
result = {}
try:
props = Document(file_path).core_properties
if props.title:
result['title'] = props.title
if props.author:
result['author'] = props.author
if props.created:
result['date'] = props.created.isoformat()
except Exception:
pass
return result
|