autochunks 0.0.8__py3-none-any.whl

autochunk/embedding/__init__.py

@@ -0,0 +1,22 @@
+
+from .base import BaseEncoder
+from .local import LocalEncoder
+from .hashing import HashingEmbedding
+from .openai import OpenAIEncoder
+from .ollama import OllamaEncoder
+
+def get_encoder(provider: str, model_name: str, **kwargs) -> BaseEncoder:
+    """
+     Factory for chunking-aware embeddings.
+    Easily extendable to TEI, OpenAI, etc.
+    """
+    if provider == "local":
+        return LocalEncoder(model_name_or_path=model_name, **kwargs)
+    elif provider == "hashing":
+        return HashingEmbedding(dim=kwargs.get("dim", 256))
+    elif provider == "openai":
+        return OpenAIEncoder(model_name=model_name, **kwargs)
+    elif provider == "ollama":
+        return OllamaEncoder(model_name=model_name, **kwargs)
+    else:
+        raise ValueError(f"Unknown embedding provider: {provider}")

autochunk/embedding/adapter.py

@@ -0,0 +1,14 @@
+
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Callable, List
+
+@dataclass
+class EmbeddingFn:
+    name: str
+    dim: int
+    fn: Callable[[List[str]], List[List[float]]]
+    cost_per_1k_tokens: float = 0.0
+
+    def __call__(self, texts: List[str]):
+        return self.fn(texts)

autochunk/embedding/base.py

@@ -0,0 +1,33 @@
+
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+class BaseEncoder(ABC):
+    """
+     interface for all AutoChunks embedding providers.
+    Designed for high-throughput batching and pluggable backends (Local, TEI, OpenAI).
+    """
+    
+    @abstractmethod
+    def embed_batch(self, texts: List[str]) -> List[List[float]]:
+        """
+        Embed a list of strings into a list of vectors.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def dimension(self) -> int:
+        """
+        Return the embedding dimension.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def model_name(self) -> str:
+        """
+        Return the model name/ID.
+        """
+        pass

autochunk/embedding/hashing.py

@@ -0,0 +1,42 @@
+
+from __future__ import annotations
+from typing import List
+import hashlib
+import numpy as np
+
+from .base import BaseEncoder
+
+class HashingEmbedding(BaseEncoder):
+    """Deterministic, offline-safe feature hashing embedding.
+    Not semantically meaningful but good for plumbing tests.
+    """
+    def __init__(self, dim: int = 256):
+        self._dim = dim
+
+    @property
+    def dimension(self) -> int:
+        return self._dim
+
+    @property
+    def model_name(self) -> str:
+        return "deterministic_hashing"
+
+    def _tok_hash(self, tok: str) -> int:
+        return int(hashlib.md5(tok.encode('utf-8')).hexdigest(), 16)
+
+    def embed_batch(self, texts: List[str]) -> List[List[float]]:
+        vecs = []
+        D = self._dim
+        for t in texts:
+            v = np.zeros(D, dtype=np.float32)
+            for tok in t.lower().split():
+                h = self._tok_hash(tok)
+                idx = h % D
+                sign = 1.0 if (h >> 1) & 1 else -1.0
+                v[idx] += sign
+            # L2 normalize
+            norm = np.linalg.norm(v)
+            if norm > 0:
+                v = v / norm
+            vecs.append(v.tolist())
+        return vecs

autochunk/embedding/local.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+from typing import List, Dict, Any, Optional
+from ..utils.logger import logger
+from .base import BaseEncoder
+
+class LocalEncoder(BaseEncoder):
+    """
+    Local Embedding Engine powered by sentence-transformers.
+    Automatically handles GPU/CPU selection and MTEB-aligned pooling logic.
+    """
+    
+    def __init__(self, model_name_or_path: str = "BAAI/bge-small-en-v1.5", device: str = None, cache_folder: str = None, trusted_orgs: List[str] = None):
+        try:
+            from sentence_transformers import SentenceTransformer
+            import torch
+        except ImportError:
+            raise ImportError("Please install sentence-transformers: pip install sentence-transformers torch")
+        
+        # Device Detection
+        if device is None:
+            if torch.cuda.is_available():
+                device = "cuda"
+            elif torch.backends.mps.is_available():
+                device = "mps"
+            else:
+                device = "cpu"
+                
+        from ..utils.logger import logger
+        logger.info(f"Using device [{device.upper()}] for embeddings.")
+        
+        if device == "cpu":
+            # Check if we are missing out on CUDA
+            try:
+                import subprocess
+                nvidia_smi = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+                if nvidia_smi.returncode == 0:
+                    logger.warning("Tip: NVIDIA GPU detected hardware-wise, but Torch is using CPU. "
+                                   "Suggest reinstalling torch with CUDA support: https://pytorch.org/get-started/locally/")
+            except Exception:
+                pass
+        
+        # Safety Check: If downloading, ensure it's from a trusted official source
+        if not cache_folder and "/" in model_name_or_path:
+            org = model_name_or_path.split("/")[0]
+            allowed = trusted_orgs or ["ds4sd", "RapidAI", "BAAI", "sentence-transformers"]
+            if org not in allowed:
+                raise ValueError(f"Security Alert: Attempting to download from untrusted source '{org}'. "
+                                 f"Trusted official orgs are: {allowed}")
+
+        self.name = model_name_or_path
+        
+        # Check if we might be downloading
+        import os
+        from ..utils.logger import logger
+        
+        # Determine the effective cache path
+        effective_cache = cache_folder or os.path.join(os.path.expanduser("~"), ".cache", "huggingface", "hub")
+        model_id_folder = "models--" + model_name_or_path.replace("/", "--")
+        is_cached = os.path.exists(os.path.join(effective_cache, model_id_folder)) or os.path.exists(model_name_or_path)
+
+        if not is_cached:
+            logger.info(f"Network Download: Local model '{model_name_or_path}' not found at {effective_cache}. Starting download from Hugging Face (Official)...")
+        else:
+            logger.info(f"Cache Hit: Using local model at {effective_cache if not os.path.exists(model_name_or_path) else model_name_or_path}")
+
+        # normalize_embeddings=True is standard for cosine similarity retrieval
+        self.model = SentenceTransformer(model_name_or_path, device=device, cache_folder=cache_folder)
+        self._dim = self.model.get_sentence_embedding_dimension()
+        
+        # Log the detected capability
+        limit = self.max_seq_length
+        logger.info(f"LocalEncoder: Loaded '{model_name_or_path}' with max sequence length: {limit} tokens (truncation @ ~{int(limit * 4 * 0.95)} chars)")
+
+    def embed_batch(self, texts: List[str]) -> List[List[float]]:
+        # Use internal LRU cache to avoid re-embedding identical text segments
+        # across multiple candidate variations (common during hyperparameter sweeps)
+        if not hasattr(self, "_cache"):
+            self._cache = {}
+        
+        results = [None] * len(texts)
+        to_embed_indices = []
+        to_embed_texts = []
+        hits = 0
+        
+        # Determine safe truncation limit
+        safe_limit = int(self.max_seq_length * 4 * 0.95)
+        
+        for i, t in enumerate(texts):
+            # Check cache first (using full text as key to avoid collisions)
+            if t in self._cache:
+                results[i] = self._cache[t]
+                hits += 1
+            else:
+                to_embed_indices.append(i)
+                # Truncate strictly for embedding model safety
+                # We store original text in cache key for recall, but embed the truncated version? 
+                # No, best to cache the truncated text result.
+                
+                # Actually, if we truncate here, we should be careful.
+                # But to save the crash, we must truncate.
+                safe_t = t[:safe_limit]
+                if len(t) > safe_limit:
+                     # Log only once to avoid spam
+                     if not hasattr(self, "_logged_truncation"):
+                         logger.warning(f"LocalEncoder: Auto-truncating input > {safe_limit} chars for model safety.")
+                         self._logged_truncation = True
+                
+                to_embed_texts.append(safe_t)
+        
+        if hits > 0:
+            logger.info(f"LocalEncoder: Cache Hits={hits}/{len(texts)}")
+        
+        if to_embed_texts:
+            import numpy as np
+            # Efficient batching is handled internally by sentence-transformers
+            embeddings = self.model.encode(to_embed_texts, show_progress_bar=False)
+            
+            # Cross-version compatibility: handle both numpy arrays and lists
+            if isinstance(embeddings, np.ndarray):
+                embeddings = embeddings.tolist()
+                
+            for relative_idx, emb in enumerate(embeddings):
+                # Map back to original global index
+                original_idx = to_embed_indices[relative_idx]
+                
+                # Get ORIGINAL text for cache key (so future calls hit cache)
+                original_text = texts[original_idx]
+                
+                self._cache[original_text] = emb
+                results[original_idx] = emb
+                
+            # Optional: Simple cache eviction if it gets too large (> 10k entries)
+            if len(self._cache) > 10000:
+                # Clear half the cache if it overflows
+                keys = list(self._cache.keys())
+                for k in keys[:5000]:
+                    del self._cache[k]
+                    
+        return results
+
+    @property
+    def dimension(self) -> int:
+        return self._dim
+
+    @property
+    def model_name(self) -> str:
+        return self.name
+
+    @property
+    def max_seq_length(self) -> int:
+        """Returns the maximum token length the model can handle."""
+        if hasattr(self.model, "max_seq_length"):
+            return self.model.max_seq_length
+        return 512  # Safe default for BERT-like models

autochunk/embedding/ollama.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+import requests
+from typing import List, Optional
+from ..utils.logger import logger
+from .base import BaseEncoder
+
+class OllamaEncoder(BaseEncoder):
+    """
+    Ollama Embedding Provider.
+    Assumes Ollama is running locally at http://localhost:11434
+    """
+    
+    def __init__(self, model_name: str = "llama3", base_url: str = "http://localhost:11434"):
+        self.name = model_name
+        self.base_url = base_url.rstrip("/")
+        self._dim = None # Will be detected on first call if possible
+
+    def embed_batch(self, texts: List[str]) -> List[List[float]]:
+        url = f"{self.base_url}/api/embed"
+        
+        embeddings = []
+        # Ollama /api/embed takes one prompt or a list
+        payload = {
+            "model": self.name,
+            "input": texts
+        }
+        
+        try:
+            response = requests.post(url, json=payload, timeout=120)
+            response.raise_for_status()
+            data = response.json()
+            
+            # Ollama returns "embeddings" which is a list of vectors
+            results = data.get("embeddings", [])
+            
+            if not results:
+                # Fallback to older /api/embeddings if /api/embed is not available or empty
+                # /api/embeddings is deprecated but sometimes still used
+                 logger.warning("Ollama /api/embed returned no results, trying sequential calls (fallback).")
+                 results = []
+                 for t in texts:
+                    res = requests.post(f"{self.base_url}/api/embeddings", json={"model": self.name, "prompt": t}, timeout=30)
+                    res.raise_for_status()
+                    results.append(res.json()["embedding"])
+            
+            if results and self._dim is None:
+                self._dim = len(results[0])
+                
+            return results
+        except Exception as e:
+            logger.error(f"Ollama Embedding failed: {e}")
+            raise
+
+    @property
+    def dimension(self) -> int:
+        if self._dim is None:
+            # Try to pulse the model to get dimension
+            try:
+                self.embed_batch(["pulsing"])
+            except:
+                return 4096 # common fallback for llama
+        return self._dim
+
+    @property
+    def model_name(self) -> str:
+        return self.name

autochunk/embedding/openai.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+import os
+import requests
+from typing import List, Optional
+from ..utils.logger import logger
+from .base import BaseEncoder
+
+class OpenAIEncoder(BaseEncoder):
+    """
+    OpenAI Embedding Provider.
+    Requires OPENAI_API_KEY environment variable.
+    """
+    
+    def __init__(self, model_name: str = "text-embedding-3-small", api_key: Optional[str] = None):
+        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
+        if not self.api_key:
+            logger.warning("OPENAI_API_KEY not found. OpenAI embeddings will fail unless provided.")
+        
+        self.name = model_name
+        # Common dimensions as fallback
+        self._dim_map = {
+            "text-embedding-3-small": 1536,
+            "text-embedding-3-large": 3072,
+            "text-embedding-ada-002": 1536
+        }
+        self._dim = self._dim_map.get(model_name, 1536)
+
+    def embed_batch(self, texts: List[str]) -> List[List[float]]:
+        if not self.api_key:
+            raise ValueError("OPENAI_API_KEY is required for OpenAI embeddings.")
+            
+        url = "https://api.openai.com/v1/embeddings"
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.api_key}"
+        }
+        
+        # OpenAI supports batching internally
+        payload = {
+            "input": texts,
+            "model": self.name
+        }
+        
+        try:
+            response = requests.post(url, headers=headers, json=payload, timeout=60)
+            response.raise_for_status()
+            data = response.json()
+            
+            # OpenAI returns data in the same order as input
+            embeddings = [item["embedding"] for item in data["data"]]
+            return embeddings
+        except Exception as e:
+            logger.error(f"OpenAI Embedding failed: {e}")
+            raise
+
+    @property
+    def dimension(self) -> int:
+        return self._dim
+
+    @property
+    def model_name(self) -> str:
+        return self.name

autochunk/embedding/tokenizer.py

@@ -0,0 +1,9 @@
+
+from dataclasses import dataclass
+from typing import Callable
+
+@dataclass
+class SimpleTokenizer:
+    name: str = "whitespace"
+    def tokens(self, text: str):
+        return text.split()

autochunk/enrichment/contextual.py

@@ -0,0 +1,29 @@
+
+from __future__ import annotations
+from typing import List, Dict, Any, Callable
+from ..utils.logger import logger
+
+class ContextualEnricher:
+    """
+    Implements Anthropic's 'Contextual Retrieval' concept.
+    Prepends a short summary of the parent document to each chunk.
+    """
+    def __init__(self, summarizer_fn: Callable[[str], str] = None):
+        self.summarizer = summarizer_fn
+
+    def enrich_batch(self, chunks: List[Dict[str, Any]], doc_text: str) -> List[Dict[str, Any]]:
+        if not self.summarizer:
+            logger.warning("No summarizer provided for ContextualEnricher. Skipping.")
+            return chunks
+        
+        try:
+            summary = self.summarizer(doc_text)
+            for chunk in chunks:
+                # Prepend the summary as context
+                original_text = chunk["text"]
+                chunk["text"] = f"[Document Summary: {summary}]\n\n{original_text}"
+                chunk["meta"]["contextual_summary"] = summary
+        except Exception as e:
+            logger.error(f"Error during contextual enrichment: {e}")
+            
+        return chunks

autochunk/eval/harness.py

@@ -0,0 +1,177 @@
+
+from __future__ import annotations
+from typing import List, Dict, Any, Optional, Callable
+import random, time
+from ..utils.text import split_sentences, whitespace_tokens
+from ..utils.hashing import content_hash
+from ..retrieval.in_memory import InMemoryIndex
+from ..eval.metrics import mrr_at_k, ndcg_at_k, recall_at_k
+from ..eval.synthetic import SyntheticQAGenerator
+from ..utils.logger import logger
+
+class EvalHarness:
+    def __init__(self, embedding_fn, k: int = 10):
+        self.embedding = embedding_fn
+        self.k = k
+        self.generator = SyntheticQAGenerator()
+
+    def build_synthetic_qa(self, docs: List[Dict], on_progress: Optional[Callable[[str], None]] = None) -> List[Dict]:
+        qa = []
+        rng = random.Random(42)
+        for d in docs:
+            sents = split_sentences(d["text"])[:20]
+            # 1. Add standard paraphrased queries
+            for s in sents[:2]:
+                query = self.generator.generate_hard_query(s, on_progress)
+                qa.append({
+                    "id": content_hash(d["id"] + query),
+                    "doc_id": d["id"],
+                    "query": query,
+                    "answer_span": s,
+                })
+            # 2. Add boundary-crossing queries (Advanced)
+            if len(sents) > 2:
+                boundary_qa = self.generator.generate_boundary_qa(d["id"], sents[:5], on_progress)
+                qa.extend(boundary_qa)
+        return qa
+
+    def evaluate(self, chunks: List[Dict], qa: List[Dict]) -> Dict[str, Any]:
+        # Build index
+        # 1. Add distractor/noise chunks to ensure search isn't too trivial
+        noise_chunks = []
+        for i in range(20):
+            noise_chunks.append({
+                "id": f"noise_{i}", 
+                "doc_id": "noise", 
+                "text": f"This is some random distractor text about something unrelated {i} to increase complexity.",
+                "meta": {}
+            })
+        
+        all_eval_chunks = chunks + noise_chunks
+        logger.info(f"EvalHarness: Encoding {len(all_eval_chunks)} chunks (including {len(noise_chunks)} noise)...")
+        
+        # Determine dynamic safety limit
+        model_limit = 512 # Fallback
+        
+        # Check if it's Hashing (which has no limit)
+        is_hashing = getattr(self.embedding, "name", "").startswith("hashing") or "HashingEmbedding" in str(type(self.embedding))
+        
+        if is_hashing:
+             MAX_CHARS = 1_000_000 # Virtually infinite
+             model_limit = 250_000
+        else:
+            if hasattr(self.embedding, "max_seq_length"):
+                model_limit = self.embedding.max_seq_length
+            elif hasattr(self.embedding, "__self__") and hasattr(self.embedding.__self__, "max_seq_length"):
+                 # Handle bound methods
+                model_limit = self.embedding.__self__.max_seq_length
+            MAX_CHARS = int(model_limit * 4 * 0.95)
+        
+        has_warned = False
+        def truncate(text: str) -> str:
+            nonlocal has_warned
+            if len(text) > MAX_CHARS:
+                if not has_warned:
+                    # Only warn if it's NOT hashing (since hashing truncation is rare/impossible with this high limit)
+                    logger.warning(f"EvalHarness: Truncating chunks > {MAX_CHARS} chars to fit embedding model ({model_limit} tokens).")
+                    has_warned = True
+                return text[:MAX_CHARS]
+            return text
+        
+        enc_start = time.time()
+        try:
+            vectors = self.embedding([truncate(c["text"]) for c in all_eval_chunks])
+        except RuntimeError as e:
+            if "expanded size" in str(e) or "512" in str(e):
+                logger.error(f"EvalHarness: Embedding failed - some chunks exceed model's max token length. Truncating aggressively...")
+                # Try with more aggressive truncation
+                vectors = self.embedding([truncate(c["text"])[:1200] for c in all_eval_chunks])
+            else:
+                raise
+        enc_time = time.time() - enc_start
+        logger.info(f"EvalHarness: Encoding complete in {enc_time:.2f}s")
+        
+        index = InMemoryIndex(dim=len(vectors[0]))
+        index.add(vectors, all_eval_chunks)
+
+        mrr, ndcg, recall, covered = 0.0, 0.0, 0.0, 0
+        
+        # --- BATCH QUERY EVALUATION ---
+        logger.info(f"EvalHarness: Encoding and searching {len(qa)} queries in batch mode...")
+        # Reuse detection logic
+        def truncate_q(text: str) -> str:
+             return truncate(text) # Use the same robust logic and warning system
+
+        query_texts = [truncate_q(item["query"]) for item in qa]
+        try:
+            query_vectors = self.embedding(query_texts)
+        except RuntimeError as e:
+            if "expanded size" in str(e):
+                logger.error(f"EvalHarness: Query embedding failed - text too long. Truncating aggressively...")
+                query_vectors = self.embedding([truncate_q(q)[:1000] for q in query_texts])
+            else:
+                raise
+        
+        # Batch search (using updated InMemoryIndex with batch support)
+        batch_hits = index.search(query_vectors, top_k=self.k)
+        
+        for i, (item, hits) in enumerate(zip(qa, batch_hits)):
+            target_doc = item["doc_id"].lower().replace("\\", "/")
+            
+            # --- TOKEN-LEVEL RECALL ---
+            answer_tokens = set(whitespace_tokens(item["answer_span"].lower()))
+            found_tokens = set()
+            
+            # --- RANKING & DCG RELEVANCE ---
+            retrieved_rels = []
+            has_perfect_match = False
+            
+            for rank, (idx, dist) in enumerate(hits):
+                c = index.meta[idx]
+                rel = 0.0
+                
+                # Check Document Match
+                if c["doc_id"].lower().replace("\\", "/") == item["doc_id"].lower().replace("\\", "/"):
+                    # Normalize whitespace for robust substring matching
+                    chunk_text_norm = " ".join(c["text"].lower().split())
+                    answer_norm = " ".join(item["answer_span"].lower().split())
+                    
+                    # 1. Full Answer Match (Highest Relevance)
+                    if answer_norm in chunk_text_norm:
+                        rel = 2.0
+                        has_perfect_match = True
+                        found_tokens.update(answer_tokens)
+                    else:
+                        # 2. Token Overlap Match (Partial Relevance)
+                        chunk_tokens = set(chunk_text_norm.split())
+                        overlap = answer_tokens.intersection(chunk_tokens)
+                        if overlap:
+                            rel = 1.0 + (len(overlap) / len(answer_tokens))
+                            found_tokens.update(overlap)
+                
+                retrieved_rels.append(rel)
+
+            # Score Aggregation
+            # Relaxed Coverage: Allow non-exact but high-overlap matches (relevance > 1.5)
+            # This accounts for markdown artifacts, minor cleaning diffs, etc.
+            if has_perfect_match or any(r > 1.5 for r in retrieved_rels):
+                covered += 1
+            
+            # MRR: Binary look (was there a perfect match in top-K?)
+            binary_rels = [1 if r >= 2.0 else 0 for r in retrieved_rels]
+            mrr += mrr_at_k(binary_rels, self.k)
+            
+            # nDCG: Uses the graduated relevance (0, 1.X, 2.0)
+            ndcg += ndcg_at_k(retrieved_rels, self.k)
+            
+            # Recall: Percentage of total answer tokens covered by all top-K results
+            current_recall = len(found_tokens) / len(answer_tokens) if answer_tokens else 0
+            recall += current_recall
+
+        n = max(1, len(qa))
+        return {
+            "mrr@k": mrr / n,
+            "ndcg@k": ndcg / n,
+            "recall@k": recall / n,
+            "coverage": covered / n,
+        }

autochunk/eval/metrics.py

@@ -0,0 +1,27 @@
+
+from __future__ import annotations
+from typing import List
+import math
+
+def dcg(rels: List[int]) -> float:
+    return sum((2**r - 1) / math.log2(i+2) for i, r in enumerate(rels))
+
+def ndcg_at_k(rels: List[int], k: int) -> float:
+    rels_k = rels[:k]
+    ideal = sorted(rels_k, reverse=True)
+    denom = dcg(ideal)
+    if denom == 0:
+        return 0.0
+    return dcg(rels_k) / denom
+
+def mrr_at_k(rels: List[int], k: int) -> float:
+    for i, r in enumerate(rels[:k]):
+        if r > 0:
+            return 1.0 / (i+1)
+    return 0.0
+
+def recall_at_k(rels: List[int], k: int, total_relevant: int) -> float:
+    if total_relevant == 0:
+        return 0.0
+    hit = sum(1 for r in rels[:k] if r > 0)
+    return hit / total_relevant