Homework Zolotarev D.R.

app/chunker.py

@@ -1,129 +1,161 @@
-"""Chunking: documents.jsonl → chunks.jsonl."""
+"""Document segmentation utilities for splitting text into manageable pieces."""
 
 import json
 from pathlib import Path
+from typing import List, Dict, Generator
+from datetime import datetime
 
-from app.config import CHUNK_MAX_CHARS, CHUNK_OVERLAP, DOCUMENTS_JSONL, CHUNKS_JSONL
+import app.config as cfg
 
+class TextSegmenter:
+    """Handles text segmentation with overlap support."""
 
-def split_paragraphs(text: str) -> list[str]:
-    """Разбивает текст на непустые абзацы."""
-    return [p.strip() for p in text.split("\n\n") if p.strip()]
+    def __init__(self, params: cfg.ProcessingParams = None):
+        self.params = params or cfg.ProcessingParams()
+        self.stats = {"total_segments": 0, "processed_docs": 0}
 
+    def _split_by_paragraphs(self, text: str) -> List[str]:
+        """Extract paragraphs from text, filtering empty ones."""
+        return [p.strip() for p in text.split("\n\n") if p.strip()]
 
-def split_long_text(text: str, max_chars: int) -> list[str]:
-    """Длинный абзац без переносов — жёсткая нарезка; overlap добавляется позже."""
-    if len(text) <= max_chars:
-        return [text]
-    parts: list[str] = []
-    start = 0
-    while start < len(text):
-        end = min(start + max_chars, len(text))
-        parts.append(text[start:end])
-        start = end
-    return parts
+    def _hard_cut(self, text: str, limit: int) -> List[str]:
+        """Hard split text when it exceeds length limit."""
+        if len(text) <= limit:
+            return [text]
 
+        pieces = []
+        for start in range(0, len(text), limit):
+            pieces.append(text[start:start + limit])
+        return pieces
 
-def apply_overlap(chunks: list[str], overlap: int, max_chars: int) -> list[str]:
-    """Добавляет overlap из предыдущего чанка в начало следующего."""
-    if overlap <= 0 or len(chunks) <= 1:
-        return chunks
-    result = [chunks[0]]
-    for i in range(1, len(chunks)):
-        prefix = chunks[i - 1][-overlap:]
-        combined = prefix + chunks[i]
-        if len(combined) > max_chars:
-            combined = prefix + chunks[i][: max_chars - len(prefix)]
-        result.append(combined)
-    return result
-
-
-def chunk_text(
-    text: str,
-    max_chars: int = CHUNK_MAX_CHARS,
-    overlap: int = CHUNK_OVERLAP,
-) -> list[str]:
-    """Нарезка по абзацам с ограничением длины и overlap между чанками."""
-    if not text.strip():
-        return []
-
-    raw_chunks: list[str] = []
-    current_parts: list[str] = []
-
-    def flush() -> None:
-        if current_parts:
-            raw_chunks.append("\n\n".join(current_parts))
-            current_parts.clear()
-
-    for para in split_paragraphs(text):
-        for piece in split_long_text(para, max_chars):
-            candidate_parts = current_parts + [piece]
-            candidate = "\n\n".join(candidate_parts)
-            if len(candidate) <= max_chars:
-                current_parts = candidate_parts
-            else:
-                flush()
-                if len(piece) <= max_chars:
-                    current_parts = [piece]
-                else:
-                    raw_chunks.extend(split_long_text(piece, max_chars))
+    def _merge_with_prev_tail(self, segments: List[str]) -> List[str]:
+        """Merge segments adding overlap from previous segment."""
+        overlap = self.params.segment_overlap
+        max_len = self.params.max_segment_size
+
+        if overlap <= 0 or len(segments) <= 1:
+            return segments
+
+        result = [segments[0]]
+        for i in range(1, len(segments)):
+            previous_tail = segments[i-1][-overlap:]
+            combined = previous_tail + segments[i]
+
+            if len(combined) > max_len:
+                combined = previous_tail + segments[i][:max_len - len(previous_tail)]
+            result.append(combined)
+
+        return result
+
+    def _segment_text(self, content: str) -> List[str]:
+        """Main segmentation pipeline for a single text."""
+        if not content or not content.strip():
+            return []
 
-    flush()
-    return apply_overlap(raw_chunks, overlap, max_chars)
+        raw_segments = []
+        current_batch = []
 
+        def flush_batch():
+            if current_batch:
+                raw_segments.append("\n\n".join(current_batch))
+                current_batch.clear()
 
-def chunk_document(doc: dict) -> list[dict]:
-    """Один документ → список чанков с метаданными."""
-    chunks = []
-    for i, text in enumerate(chunk_text(doc["text"])):
-        chunks.append(
-            {
-                "chunk_id": f"{doc['doc_id']}_{i}",
+        for paragraph in self._split_by_paragraphs(content):
+            for piece in self._hard_cut(paragraph, self.params.max_segment_size):
+                candidate = "\n\n".join(current_batch + [piece])
+
+                if len(candidate) <= self.params.max_segment_size:
+                    current_batch.append(piece)
+                else:
+                    flush_batch()
+                    if len(piece) <= self.params.max_segment_size:
+                        current_batch = [piece]
+                    else:
+                        raw_segments.extend(self._hard_cut(piece, self.params.max_segment_size))
+
+        flush_batch()
+        return self._merge_with_prev_tail(raw_segments)
+
+    def chunk_document(self, doc: Dict) -> List[Dict]:
+        """Convert document to segmented pieces with metadata."""
+        chunks = []
+        for idx, text_piece in enumerate(self._segment_text(doc.get("text", ""))):
+            chunks.append({
+                "chunk_id": f"{doc['doc_id']}_{idx}",
                 "doc_id": doc["doc_id"],
-                "name": doc["name"],
-                "text": text,
-            }
-        )
-    return chunks
-
-
-def load_documents(path: Path) -> list[dict]:
-    documents = []
-    with path.open(encoding="utf-8") as f:
-        for line in f:
-            line = line.strip()
-            if line:
-                documents.append(json.loads(line))
-    return documents
-
-
-def write_chunks(chunks: list[dict], path: Path) -> None:
-    path.parent.mkdir(parents=True, exist_ok=True)
-    with path.open("w", encoding="utf-8") as f:
-        for chunk in chunks:
-            f.write(json.dumps(chunk, ensure_ascii=False) + "\n")
-
-
-def run(
-    input_path: Path = DOCUMENTS_JSONL,
-    output_path: Path = CHUNKS_JSONL,
-) -> int:
-    if not input_path.exists():
-        raise FileNotFoundError(f"Не найден файл: {input_path}")
+                "name": doc.get("name", "Untitled"),
+                "text": text_piece,
+                "created_at": datetime.now().isoformat()
+            })
 
-    documents = load_documents(input_path)
-    all_chunks: list[dict] = []
-    for doc in documents:
-        all_chunks.extend(chunk_document(doc))
+        self.stats["processed_docs"] += 1
+        self.stats["total_segments"] += len(chunks)
+        return chunks
+
+    @staticmethod
+    def load_documents(filepath: Path) -> List[Dict]:
+        """Load documents from JSONL format."""
+        docs = []
+        if not filepath.exists():
+            return docs
+
+        with open(filepath, 'r', encoding='utf-8') as f:
+            for line in f:
+                if line.strip():
+                    docs.append(json.loads(line))
+        return docs
 
-    write_chunks(all_chunks, output_path)
-    return len(all_chunks)
+    @staticmethod
+    def save_chunks(chunks: List[Dict], output_path: Path) -> None:
+        """Save chunks to JSONL file."""
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(output_path, 'w', encoding='utf-8') as f:
+            for chunk in chunks:
+                f.write(json.dumps(chunk, ensure_ascii=False) + '\n')
 
 
-def main() -> None:
-    count = run()
-    print(f"Записано {count} чанков -> {CHUNKS_JSONL}")
+def chunk_text(text: str, max_chars: int, overlap: int) -> List[str]:
+    """Standalone function to chunk text."""
+    params = cfg.ProcessingParams(max_segment_size=max_chars, segment_overlap=overlap)
+    segmenter = TextSegmenter(params)
+    return segmenter._segment_text(text)
+
+
+def chunk_document(doc: Dict) -> List[Dict]:
+    """Standalone function to chunk a document."""
+    segmenter = TextSegmenter()
+    return segmenter.chunk_document(doc)
+
+
+def load_documents(filepath: Path) -> List[Dict]:
+    """Standalone function to load documents."""
+    return TextSegmenter.load_documents(filepath)
+
+
+def run(input_path: Path = None, output_path: Path = None) -> int:
+    """Command line entry point."""
+    if input_path is None:
+        input_path = cfg.DOCUMENTS_JSONL
+    if output_path is None:
+        output_path = cfg.CHUNKS_JSONL
+
+    segmenter = TextSegmenter()
+
+    if not input_path.exists():
+        print(f"Error: {input_path} not found. Run document parser first.")
+        return 0
+
+    documents = segmenter.load_documents(input_path)
+    all_chunks = []
+
+    for doc in documents:
+        all_chunks.extend(segmenter.chunk_document(doc))
+
+    segmenter.save_chunks(all_chunks, output_path)
+    print(f"✓ Generated {segmenter.stats['total_segments']} chunks from {segmenter.stats['processed_docs']} documents")
+    print(f"  Output: {output_path}")
+    return segmenter.stats["total_segments"]
 
 
 if __name__ == "__main__":
-    main()
+    run()

app/config.py

@@ -1,19 +1,53 @@
+"""Configuration settings for the document search system."""
+
+import os
 from pathlib import Path
+from dataclasses import dataclass
+
+# Base directory setup
+BASE_DIR = Path(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+# Data storage paths
+DATA_PATHS = {
+    "raw": BASE_DIR / "data" / "raw",
+    "processed": BASE_DIR / "data" / "processed",
+    "indexes": BASE_DIR / "data" / "index",
+}
 
-ROOT = Path(__file__).resolve().parent.parent
+# File definitions
+FILES = {
+    "source_data": DATA_PATHS["raw"] / "datasets.json",
+    "parsed_docs": DATA_PATHS["processed"] / "documents.jsonl",
+    "text_segments": DATA_PATHS["processed"] / "chunks.jsonl",
+    "vector_model": DATA_PATHS["indexes"] / "tfidf_model.pkl",
+    "term_matrix": DATA_PATHS["indexes"] / "tfidf_matrix.npz",
+    "segment_index": DATA_PATHS["indexes"] / "chunks_index.jsonl",
+}
 
-DATA_RAW = ROOT / "data" / "raw"
-DATA_PROCESSED = ROOT / "data" / "processed"
-DATA_INDEX = ROOT / "data" / "index"
+@dataclass
+class ProcessingParams:
+    """Parameters for document processing and retrieval."""
+    max_segment_size: int = 400
+    segment_overlap: int = 50
+    top_results: int = 3
+    relevance_threshold: float = 0.15
 
-RAW_DATASETS = DATA_RAW / "datasets.json"
-DOCUMENTS_JSONL = DATA_PROCESSED / "documents.jsonl"
-CHUNKS_JSONL = DATA_PROCESSED / "chunks.jsonl"
+# Initialize directories
+for path in DATA_PATHS.values():
+    path.mkdir(parents=True, exist_ok=True)
 
-VECTORIZER_PKL = DATA_INDEX / "vectorizer.pkl"
-MATRIX_NPZ = DATA_INDEX / "matrix.npz"
-INDEX_CHUNKS_JSONL = DATA_INDEX / "chunks.jsonl"
+# Aliases for easier imports
+RAW_DATASETS = FILES["source_data"]
+DOCUMENTS_JSONL = FILES["parsed_docs"]
+CHUNKS_JSONL = FILES["text_segments"]
+VECTORIZER_PKL = FILES["vector_model"]
+MATRIX_NPZ = FILES["term_matrix"]
+INDEX_CHUNKS_JSONL = FILES["segment_index"]
+DATA_INDEX = DATA_PATHS["indexes"]
 
-TOP_K = 3
-CHUNK_MAX_CHARS = 400
-CHUNK_OVERLAP = 50
+# Create a default params instance for constants
+_default_params = ProcessingParams()
+CHUNK_MAX_CHARS = _default_params.max_segment_size
+CHUNK_OVERLAP = _default_params.segment_overlap
+TOP_K = _default_params.top_results
+RELEVANCE_THRESHOLD = _default_params.relevance_threshold