ragtools 0.1.0__py3-none-any.whl

ragtools/__init__.py

@@ -0,0 +1,44 @@
+from ragtools.embedders import Embedder, SentenceTransformerEmbedder, SpladeEmbedder
+from ragtools.fusers import BordaCountFuser, Fuser, ReciprocalRankFuser
+from ragtools.generators import Generator, TransformersGenerator
+from ragtools.indexes import FaissEmbeddingIndex, Index, SparseEmbeddingIndex
+from ragtools.parsers import (
+    ChunkParser,
+    CsvFileElementParser,
+    ElementPageIndexParser,
+    ElementTextParser,
+    MdFileElementParser,
+    Parser,
+    PdfFileElementParser,
+    TextFileElementParser,
+)
+from ragtools.scorers import CrossEncoderScorer, Scorer
+from ragtools.stores import DirectoryStore, FileStore, MemoryStore, Store
+
+__all__ = [
+    "BordaCountFuser",
+    "ChunkParser",
+    "CrossEncoderScorer",
+    "CsvFileElementParser",
+    "DirectoryStore",
+    "ElementPageIndexParser",
+    "ElementTextParser",
+    "Embedder",
+    "FaissEmbeddingIndex",
+    "FileStore",
+    "Fuser",
+    "Generator",
+    "Index",
+    "MdFileElementParser",
+    "MemoryStore",
+    "Parser",
+    "PdfFileElementParser",
+    "ReciprocalRankFuser",
+    "Scorer",
+    "SentenceTransformerEmbedder",
+    "SparseEmbeddingIndex",
+    "SpladeEmbedder",
+    "Store",
+    "TextFileElementParser",
+    "TransformersGenerator",
+]

ragtools/embedders.py

@@ -0,0 +1,34 @@
+from dataclasses import dataclass
+from functools import cached_property
+from typing import Protocol, cast
+
+from sentence_transformers import SentenceTransformer, SparseEncoder
+
+
+class Embedder[E](Protocol):
+    def embed(self, chunk: str) -> E: ...
+
+
+@dataclass
+class SentenceTransformerEmbedder:
+    model_name: str
+
+    @cached_property
+    def model(self) -> SentenceTransformer:
+        return SentenceTransformer(self.model_name)
+
+    def embed(self, chunk: str) -> list[float]:
+        return self.model.encode(chunk, normalize_embeddings=True).tolist()
+
+
+@dataclass
+class SpladeEmbedder:
+    model_name: str
+
+    @cached_property
+    def model(self) -> SparseEncoder:
+        return SparseEncoder(self.model_name)
+
+    def embed(self, chunk: str) -> dict[str, float]:
+        embeddings = self.model.encode(chunk)
+        return dict(cast(list[tuple[str, float]], [*self.model.decode(embeddings)]))

ragtools/fusers.py

@@ -0,0 +1,31 @@
+from collections.abc import Hashable, Sequence
+from dataclasses import dataclass
+from typing import Protocol
+
+
+class Fuser[V](Protocol):
+    def fuse(self, *values: Sequence[V]) -> list[V]: ...
+
+
+@dataclass(slots=True)
+class ReciprocalRankFuser[V: Hashable]:
+    k: int = 60
+
+    def fuse(self, *values: Sequence[V]) -> list[V]:
+        scores: dict[V, float] = {}
+        for ranking in values:
+            for rank, value in enumerate(ranking, start=1):
+                weight = 1 / (self.k + rank)
+                scores[value] = scores.get(value, 0.0) + weight
+        return sorted(scores, key=lambda value: scores[value], reverse=True)
+
+
+@dataclass(slots=True)
+class BordaCountFuser[V: Hashable]:
+    def fuse(self, *values: Sequence[V]) -> list[V]:
+        scores: dict[V, float] = {}
+        for ranking in values:
+            for rank, value in enumerate(ranking, start=1):
+                weight = len(ranking) - rank + 1
+                scores[value] = scores.get(value, 0.0) + weight
+        return sorted(scores, key=lambda value: scores[value], reverse=True)

ragtools/generators.py

@@ -0,0 +1,28 @@
+from dataclasses import dataclass
+from functools import cached_property
+from typing import Protocol
+
+from transformers import TextGenerationPipeline, pipeline
+
+
+class Generator(Protocol):
+    def prompt(self, query: str) -> str: ...
+
+
+@dataclass
+class TransformersGenerator:
+    model_name: str
+    max_new_tokens: int
+
+    @cached_property
+    def model(self) -> TextGenerationPipeline:
+        return pipeline("text-generation", model=self.model_name)
+
+    def prompt(self, query: str) -> str:
+        chat = [{"role": "user", "content": query}]
+        outputs = self.model(
+            chat,
+            max_new_tokens=self.max_new_tokens,
+            do_sample=True,
+        )
+        return outputs[0]["generated_text"][-1]["content"]

ragtools/indexes.py

@@ -0,0 +1,153 @@
+import builtins
+from collections.abc import Hashable
+from typing import Protocol
+
+import faiss
+import numpy as np
+from scipy.sparse import csr_matrix
+
+from ragtools.stores import Store
+
+
+class Index[K, V](Store[K, V], Protocol):
+    def closest(self, key: K, k: int) -> tuple[K, ...]: ...
+
+
+class FaissEmbeddingIndex:
+    def __init__(self, dimensions: int) -> None:
+        self._index = faiss.IndexIDMap2(faiss.IndexFlatL2(dimensions))
+        self._ids: builtins.set[int] = builtins.set()
+
+    def set(self, key: int, value: list[float]) -> None:
+        self.delete(key)
+        self._index.add_with_ids(
+            np.ascontiguousarray([value], dtype=np.float32),
+            np.array([key], dtype=np.int64),
+        )
+        self._ids.add(key)
+
+    def get(self, key: int) -> list[float]:
+        if key not in self._ids:
+            raise KeyError(key)
+        return self._index.reconstruct(key).tolist()
+
+    def delete(self, key: int) -> None:
+        if key in self._ids:
+            ids = np.array([key], dtype=np.int64)
+            self._index.remove_ids(faiss.IDSelectorBatch(len(ids), faiss.swig_ptr(ids)))
+            self._ids.discard(key)
+
+    def contains(self, key: int) -> bool:
+        return key in self._ids
+
+    def keys(self) -> builtins.set[int]:
+        return builtins.set(self._ids)
+
+    def closest(self, key: int, k: int) -> tuple[int, ...]:
+        if k <= 0 or self._index.ntotal == 0:
+            return ()
+        if key not in self._ids:
+            raise KeyError(key)
+        vec = self._index.reconstruct(key)[None]
+        _, ids = self._index.search(vec, min(k + 1, self._index.ntotal))
+        return tuple(int(fid) for fid in ids[0] if fid != -1 and int(fid) != key)[:k]
+
+    def __len__(self) -> int:
+        return self._index.ntotal
+
+    def __contains__(self, key: int) -> bool:
+        return self.contains(key)
+
+
+class SparseEmbeddingIndex[K: Hashable]:
+    def __init__(self) -> None:
+        self._data = np.empty(0, dtype=np.float32)
+        self._indices = np.empty(0, dtype=np.int32)
+        self._indptr = np.zeros(1, dtype=np.int64)
+        self._rows: list[K] = []
+        self._row_of: dict[K, int] = {}
+        self._vocab: dict[str, int] = {}
+        self._tokens: list[str] = []
+
+    def set(self, key: K, value: dict[str, float]) -> None:
+        self.delete(key)
+        clean = {t: w for t, w in value.items() if w}
+        cols = []
+        for t in clean:
+            if t not in self._vocab:
+                self._vocab[t] = len(self._tokens)
+                self._tokens.append(t)
+            cols.append(self._vocab[t])
+
+        self._data = np.concatenate(
+            [self._data, np.fromiter(clean.values(), np.float32, len(clean))]
+        )
+        self._indices = np.concatenate(
+            [self._indices, np.fromiter(cols, np.int32, len(clean))]
+        )
+        self._indptr = np.append(self._indptr, len(self._data))
+        self._row_of[key] = len(self._rows)
+        self._rows.append(key)
+
+    def get(self, key: K) -> dict[str, float]:
+        if key not in self._row_of:
+            raise KeyError(key)
+        r = self._row_of[key]
+        start, end = int(self._indptr[r]), int(self._indptr[r + 1])
+        return {
+            self._tokens[c]: float(w)
+            for c, w in zip(
+                self._indices[start:end], self._data[start:end], strict=True
+            )
+        }
+
+    def delete(self, key: K) -> None:
+        if (r := self._row_of.pop(key, None)) is None:
+            return
+        start, end = int(self._indptr[r]), int(self._indptr[r + 1])
+        self._data = np.delete(self._data, slice(start, end))
+        self._indices = np.delete(self._indices, slice(start, end))
+        counts = np.diff(self._indptr)
+        keep = np.ones(counts.size, dtype=bool)
+        keep[r] = False
+        self._indptr = np.concatenate(
+            [np.zeros(1, dtype=np.int64), np.cumsum(counts[keep], dtype=np.int64)]
+        )
+        del self._rows[r]
+        self._row_of = {k: i for i, k in enumerate(self._rows)}
+
+    def contains(self, key: K) -> bool:
+        return key in self._row_of
+
+    def keys(self) -> builtins.set[K]:
+        return builtins.set(self._row_of)
+
+    def closest(self, key: K, k: int) -> tuple[K, ...]:
+        if key not in self._row_of:
+            raise KeyError(key)
+        if k <= 0 or len(self._rows) <= 1 or len(self._vocab) == 0:
+            return ()
+        r = self._row_of[key]
+        start, end = int(self._indptr[r]), int(self._indptr[r + 1])
+        q = np.zeros(len(self._vocab), dtype=np.float32)
+        q[self._indices[start:end]] = self._data[start:end]
+        scores = self._matrix @ q
+        scores[r] = -np.inf
+        scores[scores == 0.0] = -np.inf
+        k_eff = min(k, len(self._rows) - 1)
+        top = np.argpartition(-scores, k_eff - 1)[:k_eff]
+        top = top[np.argsort(-scores[top], kind="stable")]
+        return tuple(self._rows[i] for i in top if scores[i] != -np.inf)
+
+    @property
+    def _matrix(self) -> csr_matrix:
+        return csr_matrix(
+            (self._data, self._indices, self._indptr),
+            shape=(len(self._rows), max(len(self._vocab), 1)),
+        )
+
+    def __len__(self) -> int:
+        return len(self._rows)
+
+    def __contains__(self, key: K) -> bool:
+        return self.contains(key)

ragtools/parsers.py

@@ -0,0 +1,75 @@
+from dataclasses import dataclass
+from functools import cached_property
+from pathlib import Path
+from typing import Protocol, cast
+
+from semantic_chunker import TextSplitter, get_chunker
+from unstructured.documents.elements import Element
+from unstructured.partition.csv import partition_csv
+from unstructured.partition.md import partition_md
+from unstructured.partition.pdf import partition_pdf
+from unstructured.partition.text import partition_text
+
+
+class Parser[D, V](Protocol):
+    def units(self, data: D) -> V: ...
+
+
+@dataclass(slots=True)
+class CsvFileElementParser:
+    def units(self, data: Path) -> tuple[Element, ...]:
+        return tuple(partition_csv(str(data)))
+
+
+@dataclass(slots=True)
+class MdFileElementParser:
+    def units(self, data: Path) -> tuple[Element, ...]:
+        return tuple(partition_md(str(data)))
+
+
+@dataclass(slots=True)
+class PdfFileElementParser:
+    def units(self, data: Path) -> tuple[Element, ...]:
+        return tuple(partition_pdf(str(data)))
+
+
+@dataclass(slots=True)
+class TextFileElementParser:
+    def units(self, data: Path) -> tuple[Element, ...]:
+        return tuple(partition_text(str(data)))
+
+
+@dataclass(slots=True)
+class ElementTextParser:
+    def units(self, data: Element) -> str:
+        return data.text
+
+
+@dataclass(slots=True)
+class ElementPageIndexParser:
+    def units(self, data: Element) -> int:
+        page_index = data.metadata.page_number
+        return page_index if page_index is not None else -1
+
+@dataclass
+class ChunkParser:
+    model_name: str
+    chunk_size: int
+    overlap: int
+
+    @cached_property
+    def model(self) -> TextSplitter:
+        return cast(
+            "TextSplitter",
+            get_chunker(
+                self.model_name,
+                chunking_type="text",
+                tree_sitter_language=None,
+                max_tokens=self.chunk_size,
+                overlap=self.overlap,
+                trim=True,
+            ),
+        )
+
+    def units(self, data: str) -> tuple[str, ...]:
+        return tuple(self.model.chunks(data))

ragtools/scorers.py

@@ -0,0 +1,21 @@
+from dataclasses import dataclass
+from functools import cached_property
+from typing import Protocol
+
+from sentence_transformers import CrossEncoder
+
+
+class Scorer(Protocol):
+    def score(self, query: str, chunk: str) -> float: ...
+
+
+@dataclass
+class CrossEncoderScorer:
+    model_name: str
+
+    @cached_property
+    def model(self) -> CrossEncoder:
+        return CrossEncoder(self.model_name)
+
+    def score(self, query: str, chunk: str) -> float:
+        return float(self.model.predict((query, chunk)).tolist())

ragtools/stores.py

@@ -0,0 +1,118 @@
+import builtins
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Protocol
+
+import blake3
+import dill
+
+
+class Store[K, V](Protocol):
+    def set(self, key: K, value: V) -> None: ...
+    def get(self, key: K) -> V: ...
+    def delete(self, key: K) -> None: ...
+    def contains(self, key: K) -> bool: ...
+    def keys(self) -> builtins.set[K]: ...
+
+
+@dataclass(slots=True)
+class MemoryStore[K, V]:
+    items: dict[K, V]
+
+    def __init__(self) -> None:
+        self.items = {}
+
+    def set(self, key: K, value: V) -> None:
+        self.items[key] = value
+
+    def get(self, key: K) -> V:
+        return self.items[key]
+
+    def delete(self, key: K) -> None:
+        del self.items[key]
+
+    def contains(self, key: K) -> bool:
+        return key in self.items
+
+    def keys(self) -> builtins.set[K]:
+        return set(self.items.keys())
+
+
+@dataclass(slots=True)
+class FileStore[K, V]:
+    file: Path
+
+    def _load_items(self) -> dict[K, V]:
+        if not self.file.is_file():
+            return {}
+        with self.file.open("rb") as f:
+            return dill.load(f)
+
+    def _save_items(self, items: dict[K, V]) -> None:
+        self.file.parent.mkdir(parents=True, exist_ok=True)
+        with self.file.open("wb") as f:
+            dill.dump(items, f)
+
+    def set(self, key: K, value: V) -> None:
+        items = self._load_items()
+        items[key] = value
+        self._save_items(items)
+
+    def get(self, key: K) -> V:
+        items = self._load_items()
+        return items[key]
+
+    def delete(self, key: K) -> None:
+        items = self._load_items()
+        del items[key]
+        self._save_items(items)
+
+    def contains(self, key: K) -> bool:
+        items = self._load_items()
+        return key in items
+
+    def keys(self) -> builtins.set[K]:
+        items = self._load_items()
+        return set(items.keys())
+
+
+@dataclass(slots=True)
+class DirectoryStore[K, V]:
+    directory: Path
+
+    def _file_path(self, key: K) -> Path:
+        return self.directory / blake3.blake3(dill.dumps(key)).hexdigest()
+
+    def set(self, key: K, value: V) -> None:
+        path = self._file_path(key)
+        self.directory.mkdir(parents=True, exist_ok=True)
+        with path.open("wb") as f:
+            dill.dump((key, value), f)
+
+    def get(self, key: K) -> V:
+        path = self._file_path(key)
+        if not path.is_file():
+            raise KeyError(key)
+        with path.open("rb") as f:
+            _, value = dill.load(f)
+            return value
+
+    def delete(self, key: K) -> None:
+        path = self._file_path(key)
+        if not path.is_file():
+            raise KeyError(key)
+        path.unlink()
+
+    def contains(self, key: K) -> bool:
+        return self._file_path(key).is_file()
+
+    def keys(self) -> builtins.set[K]:
+        if not self.directory.is_dir():
+            return set()
+        keys_set: set[K] = set()
+        for file in self.directory.iterdir():
+            if file.is_file():
+                with file.open("rb") as f:
+                    stored_key, _ = dill.load(f)
+                    keys_set.add(stored_key)
+        return keys_set

ragtools-0.1.0.dist-info/METADATA

@@ -0,0 +1,29 @@
+Metadata-Version: 2.3
+Name: ragtools
+Version: 0.1.0
+Summary: Add your description here
+Author: Wannes Vantorre
+Author-email: Wannes Vantorre <vantorrewannes@gmail.com>
+Requires-Dist: blake3>=1.0.9
+Requires-Dist: dill>=0.4.1
+Requires-Dist: faiss-cpu>=1.15.1
+Requires-Dist: scipy>=1.18.1
+Requires-Dist: semantic-chunker>=0.2.0
+Requires-Dist: sentence-transformers>=6.0.1
+Requires-Dist: tokenizers>=0.23.2
+Requires-Dist: transformers>=5.17.0
+Requires-Dist: unstructured-inference>=1.6.13
+Requires-Dist: unstructured[csv,md,pdf]>=0.27.5
+Requires-Dist: torch>=2.13.0 ; extra == 'cpu'
+Requires-Dist: torchvision>=0.28.0 ; extra == 'cpu'
+Requires-Python: >=3.14
+Project-URL: Homepage, https://github.com/VantorreWannes/ragtools
+Project-URL: Repository, https://github.com/VantorreWannes/ragtools
+Project-URL: Issues, https://github.com/VantorreWannes/ragtools/issues
+Project-URL: Changelog, https://github.com/VantorreWannes/ragtools/releases
+Provides-Extra: cpu
+Description-Content-Type: text/markdown
+
+# ragtools
+
+Simple cache based RAG primitives

ragtools-0.1.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+ragtools/__init__.py,sha256=O0wmKZc-m8YTNlmZ9nZF3kwTM0dYcRvmYjFUmK2E3Lg,1265
+ragtools/embedders.py,sha256=PqXcJtAbivONA-xfzgA4wQsocCbjdzeGVq7ziFxgrq8,938
+ragtools/fusers.py,sha256=aOU5T0-NDmcuG3JB9HmSTeARc9b_7AvMgbW_7y4Y2y4,1104
+ragtools/generators.py,sha256=wgEuufhMSxR6Y5BqaniIzsxTUuo3FxO0AnMFURdGzoo,767
+ragtools/indexes.py,sha256=ElZE4NN7X_aesm2Vh8YNbyT2EgoGBkNUkJa7Pxa7LkE,5442
+ragtools/parsers.py,sha256=3iZ5pu6ECkZKW8MbjNd8YKUJGt0s1pmC4fJylk8_K2s,2109
+ragtools/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ragtools/scorers.py,sha256=2Wsl1DT3VJu0vVDqVke8PxwuPl2Pa9HukqxA7IJBmlA,537
+ragtools/stores.py,sha256=bg-i-f2LVCI1qfTxUYZ-N5QPC77gRF3nLP6AQh5daSw,3236
+ragtools-0.1.0.dist-info/WHEEL,sha256=-i9oRNYVXXZJUIYl5zclLIg6onEb0NLibTX34uln84w,81
+ragtools-0.1.0.dist-info/METADATA,sha256=nV9BAZ-zRnKW6nGnbRUto4Fmo_h0bv8DTTDi8z2QXrg,1046
+ragtools-0.1.0.dist-info/RECORD,,

ragtools-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.12.13
+Root-Is-Purelib: true
+Tag: py3-none-any