ragmint 0.3.1__py3-none-any.whl

ragmint/__main__.py

@@ -0,0 +1,28 @@
+from pathlib import Path
+from ragmint.tuner import RAGMint
+
+def main():
+    # Dynamically resolve the path to the installed ragmint package
+    base_dir = Path(__file__).resolve().parent
+
+    docs_path = base_dir / "experiments" / "corpus"
+    validation_file = base_dir / "experiments" / "validation_qa.json"
+
+    rag = RAGMint(
+        docs_path=str(docs_path),
+        retrievers=["faiss"],
+        embeddings=["openai/text-embedding-3-small"],
+        rerankers=["mmr"],
+    )
+
+    best, results = rag.optimize(
+        validation_set=str(validation_file),
+        metric="faithfulness",
+        search_type="bayesian",
+        trials=10,
+    )
+
+    print("Best config found:\n", best)
+
+if __name__ == "__main__":
+    main()

ragmint/autotuner.py

@@ -0,0 +1,138 @@
+"""
+Auto-RAG Tuner
+--------------
+Automatically recommends and optimizes RAG configurations based on corpus statistics.
+Integrates with RAGMint to perform full end-to-end tuning.
+"""
+
+import os
+import logging
+from statistics import mean
+from typing import Dict, Any, Tuple, List
+
+from .tuner import RAGMint
+from .core.evaluation import evaluate_config
+
+logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")
+
+
+class AutoRAGTuner:
+    def __init__(self, docs_path: str):
+        """
+        AutoRAGTuner automatically analyzes a corpus and runs an optimized RAG tuning pipeline.
+
+        Args:
+            docs_path (str): Path to the directory containing documents (.txt, .md, .rst)
+        """
+        self.docs_path = docs_path
+        self.corpus_stats = self._analyze_corpus()
+
+    # -----------------------------
+    # Corpus Analysis
+    # -----------------------------
+    def _analyze_corpus(self) -> Dict[str, Any]:
+        """Compute corpus size, average length, and number of documents."""
+        docs = []
+        total_chars = 0
+        num_docs = 0
+
+        if not os.path.exists(self.docs_path):
+            logging.warning(f"⚠️ Corpus path not found: {self.docs_path}")
+            return {"size": 0, "avg_len": 0, "num_docs": 0}
+
+        for file in os.listdir(self.docs_path):
+            if file.endswith((".txt", ".md", ".rst")):
+                with open(os.path.join(self.docs_path, file), "r", encoding="utf-8") as f:
+                    content = f.read()
+                    docs.append(content)
+                    total_chars += len(content)
+                    num_docs += 1
+
+        avg_len = int(mean([len(d) for d in docs])) if docs else 0
+        stats = {"size": total_chars, "avg_len": avg_len, "num_docs": num_docs}
+        logging.info(f"📊 Corpus stats: {stats}")
+        return stats
+
+    # -----------------------------
+    # Recommendation Logic
+    # -----------------------------
+    def recommend(self) -> Dict[str, Any]:
+        """Recommend retriever, embedding, and chunking based on corpus stats."""
+        size = self.corpus_stats.get("size", 0)
+        avg_len = self.corpus_stats.get("avg_len", 0)
+        num_docs = self.corpus_stats.get("num_docs", 0)
+
+        # Heuristic-based tuning
+        # Determine chunking heuristics first
+        if avg_len < 200:
+            chunk_size, overlap = 300, 50
+        elif avg_len < 500:
+            chunk_size, overlap = 500, 100
+        else:
+            chunk_size, overlap = 800, 150
+
+        # Determine retriever–embedding based on corpus size
+        if size <= 2000:
+            retriever = "BM25"
+            embedding_model = "sentence-transformers/all-MiniLM-L6-v2"
+        elif size <= 10000:
+            retriever = "Chroma"
+            embedding_model = "sentence-transformers/paraphrase-MiniLM-L6-v2"
+        else:
+            retriever = "FAISS"
+            embedding_model = "sentence-transformers/all-mpnet-base-v2"
+
+        strategy = "fixed" if avg_len < 400 else "sentence"
+
+        recommendation = {
+            "retriever": retriever,
+            "embedding_model": embedding_model,
+            "chunk_size": chunk_size,
+            "overlap": overlap,
+            "strategy": strategy,
+        }
+
+        logging.info(f"🔮 AutoRAG Recommendation: {recommendation}")
+        return recommendation
+
+    # -----------------------------
+    # Full Auto-Tuning
+    # -----------------------------
+    def auto_tune(
+        self,
+        validation_set: str = None,
+        metric: str = "faithfulness",
+        trials: int = 5,
+        search_type: str = "random",
+    ) -> Tuple[Dict[str, Any], List[Dict[str, Any]]]:
+        """
+        Run a full automatic optimization using RAGMint.
+
+        Automatically:
+        - Recommends initial config (retriever, embedding, chunking)
+        - Launches RAGMint optimization trials
+        - Returns best configuration and results
+        """
+        rec = self.recommend()
+
+        logging.info("🚀 Launching full AutoRAG optimization with RAGMint")
+
+        tuner = RAGMint(
+            docs_path=self.docs_path,
+            retrievers=[rec["retriever"]],
+            embeddings=[rec["embedding_model"]],
+            rerankers=["mmr"],
+            chunk_sizes=[rec["chunk_size"]],
+            overlaps=[rec["overlap"]],
+            strategies=[rec["strategy"]],
+        )
+
+        best, results = tuner.optimize(
+            validation_set=validation_set,
+            metric=metric,
+            trials=trials,
+            search_type=search_type,
+        )
+
+        logging.info(f"🏁 AutoRAG tuning complete. Best: {best}")
+        return best, results

ragmint/core/chunking.py

@@ -0,0 +1,86 @@
+from typing import List
+import re
+
+try:
+    import tiktoken
+except ImportError:
+    tiktoken = None
+
+try:
+    import nltk
+    nltk.download("punkt", quiet=True)
+    from nltk.tokenize import sent_tokenize
+except ImportError:
+    sent_tokenize = None
+
+
+class Chunker:
+    """
+    Handles text chunking strategies:
+    - fixed: character-based
+    - token: token-based (requires tiktoken)
+    - sentence: splits by full sentences (requires nltk)
+    """
+
+    def __init__(self, chunk_size: int = 500, overlap: int = 100, strategy: str = "fixed"):
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+        self.strategy = strategy
+
+    def chunk_text(self, text: str) -> List[str]:
+        """Dispatches to the correct chunking strategy."""
+        if self.strategy == "token" and tiktoken:
+            return self._chunk_by_tokens(text)
+        elif self.strategy == "sentence" and sent_tokenize:
+            return self._chunk_by_sentences(text)
+        else:
+            return self._chunk_fixed(text)
+
+    # -------------------------------
+    # Fixed-length (default)
+    # -------------------------------
+    def _chunk_fixed(self, text: str) -> List[str]:
+        chunks = []
+        start = 0
+        while start < len(text):
+            end = start + self.chunk_size
+            chunks.append(text[start:end])
+            start += self.chunk_size - self.overlap
+        return chunks
+
+    # -------------------------------
+    # Token-based (for LLM embedding)
+    # -------------------------------
+    def _chunk_by_tokens(self, text: str) -> List[str]:
+        if not tiktoken:
+            raise ImportError("tiktoken is required for token-based chunking.")
+        enc = tiktoken.get_encoding("cl100k_base")
+        tokens = enc.encode(text)
+
+        chunks = []
+        for i in range(0, len(tokens), self.chunk_size - self.overlap):
+            chunk_tokens = tokens[i:i + self.chunk_size]
+            chunks.append(enc.decode(chunk_tokens))
+        return chunks
+
+    # -------------------------------
+    # Sentence-based
+    # -------------------------------
+    def _chunk_by_sentences(self, text: str) -> List[str]:
+        if not sent_tokenize:
+            raise ImportError("nltk is required for sentence-based chunking.")
+        sentences = sent_tokenize(text)
+        chunks = []
+        current_chunk = ""
+
+        for sentence in sentences:
+            if len(current_chunk) + len(sentence) <= self.chunk_size:
+                current_chunk += " " + sentence
+            else:
+                chunks.append(current_chunk.strip())
+                current_chunk = sentence
+
+        if current_chunk:
+            chunks.append(current_chunk.strip())
+
+        return chunks

ragmint/core/embeddings.py

@@ -0,0 +1,55 @@
+import numpy as np
+from dotenv import load_dotenv
+
+try:
+    from sentence_transformers import SentenceTransformer
+except ImportError:
+    SentenceTransformer = None
+
+
+class Embeddings:
+    """
+    Wrapper for embedding backends: HuggingFace (SentenceTransformers) or Dummy.
+
+    Example:
+        model = Embeddings("huggingface", model_name="all-MiniLM-L6-v2")
+        embeddings = model.encode(["example text"])
+    """
+
+    def __init__(self, backend: str = "huggingface", model_name: str = None):
+        load_dotenv()
+        self.backend = backend.lower()
+        self.model_name = model_name or "all-MiniLM-L6-v2"
+
+        if self.backend == "huggingface":
+            if SentenceTransformer is None:
+                raise ImportError("Please install `sentence-transformers` to use HuggingFace embeddings.")
+            self.model = SentenceTransformer(self.model_name)
+            self.dim = self.model.get_sentence_embedding_dimension()
+
+        elif self.backend == "dummy":
+            self.model = None
+            self.dim = 768  # Default embedding dimension for dummy backend
+
+        else:
+            raise ValueError(f"Unsupported embedding backend: {backend}")
+
+    def encode(self, texts):
+        if isinstance(texts, str):
+            texts = [texts]
+
+        if self.backend == "huggingface":
+            embeddings = self.model.encode(texts, convert_to_numpy=True, show_progress_bar=False)
+
+        elif self.backend == "dummy":
+            # Return a NumPy array of shape (len(texts), dim)
+            embeddings = np.random.rand(len(texts), self.dim).astype(np.float32)
+
+        else:
+            raise ValueError(f"Unknown embedding backend: {self.backend}")
+
+        # ✅ Always ensure NumPy array output
+        if not isinstance(embeddings, np.ndarray):
+            embeddings = np.array(embeddings, dtype=np.float32)
+
+        return embeddings

ragmint/core/evaluation.py

@@ -0,0 +1,38 @@
+import time
+from typing import Dict, Any
+from difflib import SequenceMatcher
+
+
+class Evaluator:
+    """
+    Simple evaluation of generated answers:
+      - Faithfulness (similarity between answer and context)
+      - Latency
+    """
+
+    def __init__(self):
+        pass
+
+    def evaluate(self, query: str, answer: str, context: str) -> Dict[str, Any]:
+        start = time.time()
+        faithfulness = self._similarity(answer, context)
+        latency = time.time() - start
+
+        return {
+            "faithfulness": faithfulness,
+            "latency": latency,
+        }
+
+    def _similarity(self, a: str, b: str) -> float:
+        return SequenceMatcher(None, a, b).ratio()
+
+def evaluate_config(config, validation_data):
+    evaluator = Evaluator()
+    results = []
+    for sample in validation_data:
+        query = sample.get("query", "")
+        answer = sample.get("answer", "")
+        context = sample.get("context", "")
+        results.append(evaluator.evaluate(query, answer, context))
+    return results
+

ragmint/core/pipeline.py

@@ -0,0 +1,62 @@
+from typing import Any, Dict, Optional
+from .retriever import Retriever
+from .reranker import Reranker
+from .evaluation import Evaluator
+from .chunking import Chunker
+
+
+class RAGPipeline:
+    """
+    Core Retrieval-Augmented Generation pipeline.
+    Retrieves, reranks, and evaluates a query given the configured backends.
+    Supports text chunking for optimal retrieval performance.
+    """
+
+    def __init__(
+        self,
+        retriever: Retriever,
+        reranker: Reranker,
+        evaluator: Evaluator,
+        chunk_size: int = 500,
+        overlap: int = 100,
+        chunking_strategy: str = "fixed"
+    ):
+        self.retriever = retriever
+        self.reranker = reranker
+        self.evaluator = evaluator
+
+        # Initialize chunker for preprocessing
+        self.chunker = Chunker(chunk_size=chunk_size, overlap=overlap, strategy=chunking_strategy)
+
+    def preprocess_docs(self, documents):
+        """Applies the selected chunking strategy to the document set."""
+        all_chunks = []
+        for doc in documents:
+            chunks = self.chunker.chunk_text(doc)
+            all_chunks.extend(chunks)
+        return all_chunks
+
+    def run(self, query: str, top_k: int = 5, use_chunking: bool = True) -> Dict[str, Any]:
+        # Optional preprocessing step
+        if use_chunking and hasattr(self.retriever, "documents") and self.retriever.documents:
+            self.retriever.documents = self.preprocess_docs(self.retriever.documents)
+
+        # Retrieve documents
+        retrieved_docs = self.retriever.retrieve(query, top_k=top_k)
+
+        # Rerank
+        reranked_docs = self.reranker.rerank(query, retrieved_docs)
+
+        # Construct pseudo-answer
+        answer = reranked_docs[0]["text"] if reranked_docs else ""
+        context = "\n".join([d["text"] for d in reranked_docs])
+
+        # Evaluate
+        metrics = self.evaluator.evaluate(query, answer, context)
+
+        return {
+            "query": query,
+            "answer": answer,
+            "docs": reranked_docs,
+            "metrics": metrics,
+        }

ragmint/core/reranker.py

@@ -0,0 +1,62 @@
+from typing import List, Dict, Any
+import numpy as np
+
+
+class Reranker:
+    """
+    Supports:
+      - MMR (Maximal Marginal Relevance)
+      - Dummy CrossEncoder (for demonstration)
+    """
+
+    def __init__(self, mode: str = "mmr", lambda_param: float = 0.5, seed: int = 42):
+        self.mode = mode
+        self.lambda_param = lambda_param
+        np.random.seed(seed)
+
+    def rerank(self, query: str, docs: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        if not docs:
+            return []
+
+        if self.mode == "crossencoder":
+            return self._crossencoder_rerank(query, docs)
+        return self._mmr_rerank(query, docs)
+
+    def _mmr_rerank(self, query: str, docs: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Perform MMR reranking using dummy similarity scores."""
+        selected = []
+        remaining = docs.copy()
+
+        while remaining and len(selected) < len(docs):
+            if not selected:
+                # pick doc with highest base score
+                best = max(remaining, key=lambda d: d["score"])
+            else:
+                # MMR balancing between relevance and diversity
+                mmr_scores = []
+                for d in remaining:
+                    max_div = max(
+                        [self._similarity(d["text"], s["text"]) for s in selected],
+                        default=0,
+                    )
+                    mmr_score = (
+                        self.lambda_param * d["score"]
+                        - (1 - self.lambda_param) * max_div
+                    )
+                    mmr_scores.append(mmr_score)
+                best = remaining[int(np.argmax(mmr_scores))]
+            selected.append(best)
+            remaining.remove(best)
+
+        return selected
+
+    def _crossencoder_rerank(self, query: str, docs: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Adds a small random perturbation to simulate crossencoder reranking."""
+        for d in docs:
+            d["score"] += np.random.uniform(0, 0.1)
+        return sorted(docs, key=lambda d: d["score"], reverse=True)
+
+    def _similarity(self, a: str, b: str) -> float:
+        """Dummy similarity function between two strings."""
+        # Deterministic pseudo-similarity based on hash
+        return abs(hash(a + b)) % 100 / 100.0

ragmint/core/retriever.py

@@ -0,0 +1,165 @@
+from typing import List, Dict, Any, Optional
+import numpy as np
+from .embeddings import Embeddings
+
+# Optional imports
+try:
+    import faiss
+except ImportError:
+    faiss = None
+
+try:
+    import chromadb
+except ImportError:
+    chromadb = None
+
+try:
+    from sklearn.neighbors import BallTree
+except ImportError:
+    BallTree = None
+
+try:
+    from rank_bm25 import BM25Okapi
+except ImportError:
+    BM25Okapi = None
+
+
+class Retriever:
+    """
+    Multi-backend retriever supporting:
+        - "numpy"  : basic cosine similarity (dense)
+        - "faiss"  : high-performance dense retriever
+        - "chroma" : persistent vector DB
+        - "sklearn": BallTree (cosine or Euclidean)
+        - "bm25"   : lexical retriever using Rank-BM25
+
+    Example:
+        retriever = Retriever(embedder, documents=["A", "B", "C"], backend="bm25")
+        results = retriever.retrieve("example query", top_k=3)
+    """
+
+    def __init__(
+        self,
+        embedder: Optional[Embeddings] = None,
+        documents: Optional[List[str]] = None,
+        embeddings: Optional[np.ndarray] = None,
+        backend: str = "numpy",
+    ):
+        self.embedder = embedder
+        self.documents = documents or []
+        self.backend = backend.lower()
+        self.embeddings = None
+        self.index = None
+        self.client = None
+        self.bm25 = None
+
+        # Initialize embeddings for dense backends
+        if self.backend not in ["bm25"]:
+            if embeddings is not None:
+                self.embeddings = np.array(embeddings)
+            elif self.documents and self.embedder:
+                self.embeddings = self.embedder.encode(self.documents)
+            else:
+                self.embeddings = np.zeros((0, getattr(self.embedder, "dim", 768)))
+
+            # Normalize for cosine
+            if self.embeddings.size > 0:
+                self.embeddings = self._normalize(self.embeddings)
+
+        # Initialize backend
+        self._init_backend()
+
+    # ------------------------
+    # Backend Initialization
+    # ------------------------
+    def _init_backend(self):
+        if self.backend == "faiss":
+            if faiss is None:
+                raise ImportError("faiss not installed. Run `pip install faiss-cpu`.")
+            self.index = faiss.IndexFlatIP(self.embedder.dim)
+            self.index.add(self.embeddings.astype("float32"))
+
+        elif self.backend == "chroma":
+            if chromadb is None:
+                raise ImportError("chromadb not installed. Run `pip install chromadb`.")
+            self.client = chromadb.Client()
+            self.collection = self.client.create_collection(name="ragmint_retriever")
+            for i, doc in enumerate(self.documents):
+                self.collection.add(
+                    ids=[str(i)],
+                    documents=[doc],
+                    embeddings=[self.embeddings[i].tolist()],
+                )
+
+        elif self.backend == "sklearn":
+            if BallTree is None:
+                raise ImportError("scikit-learn not installed. Run `pip install scikit-learn`.")
+            self.index = BallTree(self.embeddings)
+
+        elif self.backend == "bm25":
+            if BM25Okapi is None:
+                raise ImportError("rank-bm25 not installed. Run `pip install rank-bm25`.")
+            tokenized_corpus = [doc.lower().split() for doc in self.documents]
+            self.bm25 = BM25Okapi(tokenized_corpus)
+
+        elif self.backend != "numpy":
+            raise ValueError(f"Unsupported retriever backend: {self.backend}")
+
+    # ------------------------
+    # Retrieval
+    # ------------------------
+    def retrieve(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
+        if len(self.documents) == 0:
+            return [{"text": "", "score": 0.0}]
+
+        # BM25 retrieval (lexical)
+        if self.backend == "bm25":
+            tokenized_query = query.lower().split()
+            scores = self.bm25.get_scores(tokenized_query)
+            top_indices = np.argsort(scores)[::-1][:top_k]
+            return [
+                {"text": self.documents[i], "score": float(scores[i])}
+                for i in top_indices
+            ]
+
+        # Dense retrieval (others)
+        if self.embeddings is None or self.embeddings.size == 0:
+            return [{"text": "", "score": 0.0}]
+
+        query_vec = self.embedder.encode([query])[0]
+        query_vec = self._normalize(query_vec)
+
+        if self.backend == "numpy":
+            scores = np.dot(self.embeddings, query_vec)
+            top_indices = np.argsort(scores)[::-1][:top_k]
+            return [{"text": self.documents[i], "score": float(scores[i])} for i in top_indices]
+
+        elif self.backend == "faiss":
+            query_vec = np.expand_dims(query_vec.astype("float32"), axis=0)
+            scores, indices = self.index.search(query_vec, top_k)
+            return [{"text": self.documents[int(i)], "score": float(scores[0][j])} for j, i in enumerate(indices[0])]
+
+        elif self.backend == "chroma":
+            results = self.collection.query(query_texts=[query], n_results=top_k)
+            docs = results["documents"][0]
+            scores = results["distances"][0]
+            return [{"text": d, "score": 1 - s} for d, s in zip(docs, scores)]
+
+        elif self.backend == "sklearn":
+            distances, indices = self.index.query([query_vec], k=top_k)
+            scores = 1 - distances[0]
+            return [{"text": self.documents[int(i)], "score": float(scores[j])} for j, i in enumerate(indices[0])]
+
+        else:
+            raise ValueError(f"Unknown backend: {self.backend}")
+
+    # ------------------------
+    # Utils
+    # ------------------------
+    @staticmethod
+    def _normalize(vectors: np.ndarray) -> np.ndarray:
+        if vectors.ndim == 1:
+            norm = np.linalg.norm(vectors)
+            return vectors / norm if norm > 0 else vectors
+        norms = np.linalg.norm(vectors, axis=1, keepdims=True)
+        return np.divide(vectors, norms, out=np.zeros_like(vectors), where=norms != 0)

ragmint/experiments/validation_qa.json

@@ -0,0 +1,14 @@
+[
+  {
+    "query": "What is Retrieval-Augmented Generation?",
+    "expected_answer": "A technique that combines information retrieval with language generation to improve factual accuracy."
+  },
+  {
+    "query": "What is the role of embeddings in a RAG system?",
+    "expected_answer": "They represent text as numerical vectors for similarity-based retrieval."
+  },
+  {
+    "query": "What is Maximal Marginal Relevance used for?",
+    "expected_answer": "To select diverse and relevant documents during reranking."
+  }
+]