File size: 7,860 Bytes
aca8ab4 |
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 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 |
"""
PDF processing and text extraction with chunking.
"""
import logging
from pathlib import Path
from typing import List, Optional
import hashlib
import tiktoken
from pypdf import PdfReader
from utils.schemas import PaperChunk, Paper
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
class PDFProcessor:
"""Process PDFs and extract text with intelligent chunking."""
def __init__(
self,
chunk_size: int = 500,
chunk_overlap: int = 50,
encoding_name: str = "cl100k_base"
):
"""
Initialize PDF processor.
Args:
chunk_size: Target chunk size in tokens
chunk_overlap: Overlap between chunks in tokens
encoding_name: Tiktoken encoding name
"""
self.chunk_size = chunk_size
self.chunk_overlap = chunk_overlap
self.encoding = tiktoken.get_encoding(encoding_name)
def extract_text(self, pdf_path: Path) -> Optional[str]:
"""
Extract text from PDF.
Args:
pdf_path: Path to PDF file
Returns:
Extracted text or None if extraction fails
"""
try:
reader = PdfReader(str(pdf_path))
text_parts = []
for page_num, page in enumerate(reader.pages, start=1):
try:
text = page.extract_text()
if text.strip():
text_parts.append(f"[Page {page_num}]\n{text}")
except Exception as e:
logger.warning(f"Failed to extract text from page {page_num}: {str(e)}")
continue
if not text_parts:
logger.error(f"No text extracted from {pdf_path}")
return None
full_text = "\n\n".join(text_parts)
logger.info(f"Extracted {len(full_text)} characters from {pdf_path.name}")
return full_text
except Exception as e:
logger.error(f"Error reading PDF {pdf_path}: {str(e)}")
return None
def _generate_chunk_id(self, paper_id: str, chunk_index: int) -> str:
"""Generate unique chunk ID."""
content = f"{paper_id}_{chunk_index}"
return hashlib.md5(content.encode()).hexdigest()
def chunk_text(
self,
text: str,
paper: Paper
) -> List[PaperChunk]:
"""
Chunk text into overlapping segments.
Args:
text: Full text to chunk
paper: Paper metadata
Returns:
List of PaperChunk objects
"""
chunks = []
tokens = self.encoding.encode(text)
# Extract page information from text
page_markers = []
lines = text.split('\n')
current_char = 0
for line in lines:
if line.startswith('[Page ') and line.endswith(']'):
try:
page_num = int(line[6:-1])
page_markers.append((current_char, page_num))
except ValueError:
pass
current_char += len(line) + 1
chunk_index = 0
start_idx = 0
while start_idx < len(tokens):
# Calculate end index
end_idx = min(start_idx + self.chunk_size, len(tokens))
# Get chunk tokens and decode
chunk_tokens = tokens[start_idx:end_idx]
chunk_text = self.encoding.decode(chunk_tokens)
# Determine page number
chunk_start_char = len(self.encoding.decode(tokens[:start_idx]))
page_number = self._get_page_number(chunk_start_char, page_markers)
# Extract section if possible (simple heuristic)
section = self._extract_section(chunk_text)
# Create metadata with type validation
try:
# Ensure authors is a list of strings
authors_metadata = paper.authors
if not isinstance(authors_metadata, list):
logger.warning(f"Paper {paper.arxiv_id} has invalid authors type: {type(authors_metadata)}, converting to list")
authors_metadata = [str(authors_metadata)] if authors_metadata else []
# Ensure title is a string
title_metadata = str(paper.title) if paper.title else ""
metadata = {
"title": title_metadata,
"authors": authors_metadata,
"chunk_index": chunk_index,
"token_count": len(chunk_tokens)
}
except Exception as e:
logger.warning(f"Error creating metadata for chunk {chunk_index}: {str(e)}, using fallback")
metadata = {
"title": str(paper.title) if hasattr(paper, 'title') else "",
"authors": [],
"chunk_index": chunk_index,
"token_count": len(chunk_tokens)
}
# Create chunk with validated data
try:
chunk = PaperChunk(
chunk_id=self._generate_chunk_id(paper.arxiv_id, chunk_index),
paper_id=paper.arxiv_id,
content=chunk_text.strip(),
section=section,
page_number=page_number,
arxiv_url=str(paper.pdf_url) if paper.pdf_url else "",
metadata=metadata
)
chunks.append(chunk)
except Exception as e:
logger.error(f"Error creating chunk {chunk_index} for paper {paper.arxiv_id}: {str(e)}")
# Continue processing other chunks rather than failing completely
continue
# Move to next chunk with overlap
start_idx += self.chunk_size - self.chunk_overlap
chunk_index += 1
logger.info(f"Created {len(chunks)} chunks for paper {paper.arxiv_id}")
return chunks
def _get_page_number(
self,
char_position: int,
page_markers: List[tuple]
) -> Optional[int]:
"""Determine page number for character position."""
if not page_markers:
return None
for i, (marker_pos, page_num) in enumerate(page_markers):
if char_position < marker_pos:
return page_markers[i - 1][1] if i > 0 else None
return page_markers[-1][1]
def _extract_section(self, text: str) -> Optional[str]:
"""
Extract section name from chunk (simple heuristic).
Looks for common section headers.
"""
section_keywords = [
'abstract', 'introduction', 'related work', 'methodology',
'method', 'experiments', 'results', 'discussion',
'conclusion', 'references', 'appendix'
]
lines = text.split('\n')[:5] # Check first 5 lines
for line in lines:
line_lower = line.lower().strip()
for keyword in section_keywords:
if keyword in line_lower and len(line.split()) < 10:
return line.strip()
return None
def process_paper(
self,
pdf_path: Path,
paper: Paper
) -> List[PaperChunk]:
"""
Process a paper PDF into chunks.
Args:
pdf_path: Path to PDF file
paper: Paper metadata
Returns:
List of PaperChunk objects
"""
# Extract text
text = self.extract_text(pdf_path)
if not text:
logger.error(f"Failed to extract text from {pdf_path}")
return []
# Chunk text
chunks = self.chunk_text(text, paper)
return chunks
|