Spaces:
Sleeping
Sleeping
| import os | |
| import re | |
| from typing import List, Dict, Any | |
| from pypdf import PdfReader | |
| def extract_text_from_pdf(pdf_path: str) -> List[Dict[str, Any]]: | |
| """ | |
| Extracts text from a PDF file page by page. | |
| """ | |
| if not os.path.exists(pdf_path): | |
| raise FileNotFoundError(f"PDF file not found at: {pdf_path}") | |
| pages_data = [] | |
| try: | |
| reader = PdfReader(pdf_path) | |
| num_pages = len(reader.pages) | |
| if num_pages == 0: | |
| raise ValueError(f"The PDF file '{pdf_path}' has 0 pages or could not be read.") | |
| for i in range(num_pages): | |
| page = reader.pages[i] | |
| text = page.extract_text() | |
| if text: | |
| cleaned_text = re.sub(r'\s+', ' ', text).strip() | |
| if cleaned_text: | |
| pages_data.append({ | |
| "page_num": i + 1, | |
| "text": cleaned_text | |
| }) | |
| except Exception as e: | |
| if isinstance(e, ValueError) or isinstance(e, FileNotFoundError): | |
| raise e | |
| raise Exception(f"Failed to read PDF file '{pdf_path}': {str(e)}") | |
| if not pages_data: | |
| raise ValueError(f"No readable text could be extracted from PDF file '{pdf_path}'.") | |
| return pages_data | |
| def extract_text_from_docx(docx_path: str) -> List[Dict[str, Any]]: | |
| """ | |
| Extracts text from a Word DOCX file paragraph by paragraph. | |
| """ | |
| if not os.path.exists(docx_path): | |
| raise FileNotFoundError(f"DOCX file not found at: {docx_path}") | |
| try: | |
| import docx | |
| except ImportError: | |
| raise ImportError("The 'python-docx' library is required to parse DOCX files. " | |
| "Please install it using 'pip install python-docx'.") | |
| try: | |
| doc = docx.Document(docx_path) | |
| paragraphs_text = [] | |
| for para in doc.paragraphs: | |
| clean_p = para.text.strip() | |
| if clean_p: | |
| paragraphs_text.append(clean_p) | |
| # Merge paragraphs with double newlines | |
| full_text = "\n\n".join(paragraphs_text) | |
| if not full_text.strip(): | |
| raise ValueError(f"The DOCX file '{docx_path}' has no readable text content.") | |
| return [{ | |
| # DOCX does not have native easy page-numbering without layout parsing, | |
| # so we treat it as page_num = 1 | |
| "page_num": 1, | |
| "text": full_text | |
| }] | |
| except Exception as e: | |
| if isinstance(e, ValueError) or isinstance(e, FileNotFoundError) or isinstance(e, ImportError): | |
| raise e | |
| raise Exception(f"Failed to read DOCX file '{docx_path}': {str(e)}") | |
| def extract_text_from_txt(txt_path: str) -> List[Dict[str, Any]]: | |
| """ | |
| Extracts text from a plain TXT file. | |
| """ | |
| if not os.path.exists(txt_path): | |
| raise FileNotFoundError(f"TXT file not found at: {txt_path}") | |
| try: | |
| with open(txt_path, 'r', encoding='utf-8', errors='ignore') as f: | |
| text = f.read() | |
| cleaned_text = text.strip() | |
| if not cleaned_text: | |
| raise ValueError(f"The TXT file '{txt_path}' is empty.") | |
| return [{ | |
| "page_num": 1, | |
| "text": cleaned_text | |
| }] | |
| except Exception as e: | |
| if isinstance(e, ValueError) or isinstance(e, FileNotFoundError): | |
| raise e | |
| raise Exception(f"Failed to read TXT file '{txt_path}': {str(e)}") | |
| def extract_text_from_file(file_path: str) -> List[Dict[str, Any]]: | |
| """ | |
| Dispatcher to extract text from a file based on its file extension. | |
| """ | |
| ext = os.path.splitext(file_path)[1].lower() | |
| if ext == '.pdf': | |
| return extract_text_from_pdf(file_path) | |
| elif ext == '.docx': | |
| return extract_text_from_docx(file_path) | |
| elif ext in ('.txt', '.md', '.py', '.json', '.env', '.csv'): | |
| return extract_text_from_txt(file_path) | |
| else: | |
| raise ValueError(f"Unsupported file type '{ext}'. Supported types: .pdf, .docx, .txt (and code files like .md, .py)") | |
| def chunk_text(pages_data: List[Dict[str, Any]], chunk_size: int, chunk_overlap: int, source_name: str) -> List[Dict[str, Any]]: | |
| """ | |
| Splits page/document texts into chunks using a character-based sliding window with overlap. | |
| Args: | |
| pages_data: List of dictionaries with 'page_num' and 'text'. | |
| chunk_size: Target size of each chunk in characters. | |
| chunk_overlap: Overlap between consecutive chunks in characters. | |
| source_name: Base name of the source file. | |
| Returns: | |
| List of chunks with source metadata. | |
| """ | |
| if chunk_size <= 0: | |
| raise ValueError("chunk_size must be a positive integer.") | |
| if chunk_overlap < 0: | |
| raise ValueError("chunk_overlap must be non-negative.") | |
| if chunk_overlap >= chunk_size: | |
| raise ValueError("chunk_overlap must be strictly less than chunk_size.") | |
| chunks = [] | |
| global_chunk_idx = 1 | |
| for page in pages_data: | |
| text = page["text"] | |
| page_num = page["page_num"] | |
| if len(text) <= chunk_size: | |
| chunks.append({ | |
| "text": text, | |
| "page_num": page_num, | |
| "chunk_index": global_chunk_idx, | |
| "source_name": source_name | |
| }) | |
| global_chunk_idx += 1 | |
| continue | |
| start = 0 | |
| while start < len(text): | |
| end = start + chunk_size | |
| if end >= len(text): | |
| chunk_text_slice = text[start:] | |
| chunks.append({ | |
| "text": chunk_text_slice.strip(), | |
| "page_num": page_num, | |
| "chunk_index": global_chunk_idx, | |
| "source_name": source_name | |
| }) | |
| global_chunk_idx += 1 | |
| break | |
| search_limit = min(chunk_overlap, 50) | |
| split_at = end | |
| for offset in range(search_limit): | |
| if text[end - offset].isspace(): | |
| split_at = end - offset | |
| break | |
| chunk_text_slice = text[start:split_at].strip() | |
| if chunk_text_slice: | |
| chunks.append({ | |
| "text": chunk_text_slice, | |
| "page_num": page_num, | |
| "chunk_index": global_chunk_idx, | |
| "source_name": source_name | |
| }) | |
| global_chunk_idx += 1 | |
| next_start = split_at - chunk_overlap | |
| if next_start <= start: | |
| start = split_at | |
| else: | |
| start = next_start | |
| return chunks | |