Ollama RAG: Build a Private Retrieval Pipeline That Never Leaves Your Machine

# Install Ollama (macOS, Linux, Windows)
curl -fsSL https://ollama.com/install.sh | sh

# Pull the embedding model
ollama pull nomic-embed-text

# Pull a generation model
ollama pull llama3.1

# Verify both models are available
ollama list

# Generate embeddings (batch)
POST http://localhost:11434/api/embed
{
  "model": "nomic-embed-text",
  "input": ["First document chunk", "Second document chunk"]
}
# Returns: { "embeddings": [[0.123, -0.456, ...], [...]] }

# Chat completion (generation)
POST http://localhost:11434/api/chat
{
  "model": "llama3.1",
  "messages": [
    {"role": "system", "content": "Answer based on the provided context."},
    {"role": "user", "content": "Context: ...\n\nQuestion: ..."}
  ],
  "stream": false
}

import os
from pathlib import Path

def load_documents(directory: str) -> list[dict]:
    """Load text files from a directory."""
    docs = []
    supported = {".txt", ".md", ".py", ".ts", ".js", ".html", ".csv"}

    for path in Path(directory).rglob("*"):
        if path.suffix in supported and path.is_file():
            text = path.read_text(encoding="utf-8", errors="ignore")
            docs.append({
                "text": text,
                "metadata": {
                    "source": str(path),
                    "filename": path.name,
                    "extension": path.suffix,
                }
            })
    return docs

# For PDFs, add PyPDF:
# pip install pypdf
from pypdf import PdfReader

def load_pdf(path: str) -> list[dict]:
    reader = PdfReader(path)
    docs = []
    for i, page in enumerate(reader.pages):
        text = page.extract_text() or ""
        if text.strip():
            docs.append({
                "text": text,
                "metadata": {"source": path, "page": i + 1}
            })
    return docs

def split_text(
    text: str,
    chunk_size: int = 800,
    chunk_overlap: int = 150,
    separators: list[str] | None = None,
) -> list[str]:
    """
    Recursively split text, preferring natural boundaries.
    Tries separators in order: paragraph > newline > sentence > space.
    """
    if separators is None:
        separators = ["\n\n", "\n", ". ", " "]

    chunks = []
    current_sep = separators[0]

    # If text fits in one chunk, return it
    if len(text) <= chunk_size:
        return [text.strip()] if text.strip() else []

    # Split on current separator
    parts = text.split(current_sep)

    current_chunk = ""
    for part in parts:
        candidate = current_chunk + current_sep + part if current_chunk else part

        if len(candidate) <= chunk_size:
            current_chunk = candidate
        else:
            if current_chunk:
                chunks.append(current_chunk.strip())
            # If a single part exceeds chunk_size, try next separator
            if len(part) > chunk_size and len(separators) > 1:
                chunks.extend(
                    split_text(part, chunk_size, chunk_overlap, separators[1:])
                )
                current_chunk = ""
            else:
                current_chunk = part

    if current_chunk.strip():
        chunks.append(current_chunk.strip())

    # Add overlap: prepend tail of previous chunk to each chunk
    if chunk_overlap > 0 and len(chunks) > 1:
        overlapped = [chunks[0]]
        for i in range(1, len(chunks)):
            prev_tail = chunks[i - 1][-chunk_overlap:]
            overlapped.append(prev_tail + " " + chunks[i])
        chunks = overlapped

    return chunks

import requests
from typing import List

OLLAMA_URL = "http://localhost:11434"

def embed_texts(texts: list[str], model: str = "nomic-embed-text") -> list[list[float]]:
    """Embed a batch of texts using Ollama's /api/embed endpoint."""
    response = requests.post(
        f"{OLLAMA_URL}/api/embed",
        json={"model": model, "input": texts},
    )
    response.raise_for_status()
    return response.json()["embeddings"]

# Example: embed 100 chunks in one call
chunks = ["chunk 1 text...", "chunk 2 text...", ...]  # your split documents
vectors = embed_texts(chunks)
print(f"Embedded {len(vectors)} chunks, each {len(vectors[0])} dimensions")

# pip install chromadb
import chromadb

# Create a persistent client (data survives restarts)
client = chromadb.PersistentClient(path="./chroma_db")

# Create or get a collection
collection = client.get_or_create_collection(
    name="my-documents",
    metadata={"hnsw:space": "cosine"},  # cosine similarity
)

def ingest_documents(docs: list[dict], collection):
    """Split, embed, and store documents in ChromaDB."""
    all_chunks = []
    all_metadatas = []
    all_ids = []

    for doc in docs:
        chunks = split_text(doc["text"])
        for i, chunk in enumerate(chunks):
            all_chunks.append(chunk)
            all_metadatas.append({
                **doc["metadata"],
                "chunk_index": i,
            })
            all_ids.append(f"{doc['metadata']['filename']}_{i}")

    # Embed all chunks
    embeddings = embed_texts(all_chunks)

    # Store in ChromaDB (batch for large corpora)
    batch_size = 500
    for start in range(0, len(all_chunks), batch_size):
        end = start + batch_size
        collection.add(
            ids=all_ids[start:end],
            embeddings=embeddings[start:end],
            documents=all_chunks[start:end],
            metadatas=all_metadatas[start:end],
        )

    print(f"Ingested {len(all_chunks)} chunks from {len(docs)} documents")

# Run ingestion
docs = load_documents("./my_docs")
ingest_documents(docs, collection)

import requests

def query_rag(
    question: str,
    collection,
    n_results: int = 5,
    model: str = "llama3.1",
) -> str:
    """Full RAG query: embed question, retrieve, generate."""

    # Step 1: Embed the question
    q_embedding = embed_texts([question])[0]

    # Step 2: Retrieve similar chunks
    results = collection.query(
        query_embeddings=[q_embedding],
        n_results=n_results,
    )

    # Step 3: Build context from retrieved chunks
    context_parts = []
    for doc, meta in zip(results["documents"][0], results["metadatas"][0]):
        source = meta.get("source", "unknown")
        context_parts.append(f"[Source: {source}]\n{doc}")

    context = "\n\n---\n\n".join(context_parts)

    # Step 4: Generate answer with Ollama
    prompt = f"""Use the following context to answer the question.
If the context does not contain enough information, say so.
Do not make up information.

Context:
{context}

Question: {question}"""

    response = requests.post(
        f"{OLLAMA_URL}/api/chat",
        json={
            "model": model,
            "messages": [
                {"role": "system", "content": "You answer questions based on provided context. Be precise and cite sources."},
                {"role": "user", "content": prompt},
            ],
            "stream": False,
            "options": {"num_ctx": 8192},
        },
    )
    response.raise_for_status()
    return response.json()["message"]["content"]

# Usage
answer = query_rag("What is the refund policy?", collection)
print(answer)

# Increase context window (costs ~1 GB VRAM per 4k tokens)
# Set in the API call:
"options": {"num_ctx": 8192}

# Or create a Modelfile with persistent settings:
# Modelfile.rag
FROM llama3.1
PARAMETER num_ctx 8192
PARAMETER num_gpu 999          # offload all layers to GPU
PARAMETER num_thread 8         # CPU threads for non-GPU work

# Build the custom model
ollama create llama3.1-rag -f Modelfile.rag

# For Apple Silicon: Flash Attention is enabled automatically
# For CUDA: Flash Attention enabled by default since Ollama 0.3+

# Monitor VRAM usage during queries
# macOS: Activity Monitor > Memory (Wired + Compressed)
# Linux: nvidia-smi -l 1

# pip install rank-bm25
from rank_bm25 import BM25Okapi
import numpy as np

class HybridRetriever:
    def __init__(self, collection, documents: list[str]):
        self.collection = collection
        self.documents = documents

        # Build BM25 index from the same documents
        tokenized = [doc.lower().split() for doc in documents]
        self.bm25 = BM25Okapi(tokenized)

    def search(
        self,
        query: str,
        n_results: int = 5,
        vector_weight: float = 0.7,
    ) -> list[dict]:
        # Vector search
        q_embedding = embed_texts([query])[0]
        vector_results = self.collection.query(
            query_embeddings=[q_embedding],
            n_results=n_results * 2,  # over-retrieve
        )

        # BM25 search
        tokenized_query = query.lower().split()
        bm25_scores = self.bm25.get_scores(tokenized_query)
        bm25_top = np.argsort(bm25_scores)[::-1][:n_results * 2]

        # Normalize and combine scores
        vector_ids = set(vector_results["ids"][0])
        combined = {}

        for i, doc_id in enumerate(vector_results["ids"][0]):
            # Vector score: higher rank = higher score
            combined[doc_id] = vector_weight * (1.0 - i / len(vector_results["ids"][0]))

        for rank, idx in enumerate(bm25_top):
            doc_id = f"doc_{idx}"  # match your ID scheme
            kw_score = (1 - vector_weight) * (1.0 - rank / len(bm25_top))
            combined[doc_id] = combined.get(doc_id, 0) + kw_score

        # Sort by combined score and return top-n
        ranked = sorted(combined.items(), key=lambda x: x[1], reverse=True)
        return ranked[:n_results]

# pip install sentence-transformers
from sentence_transformers import CrossEncoder

# Load a small cross-encoder (runs on CPU in ~50ms per pair)
reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

def rerank(query: str, documents: list[str], top_k: int = 5) -> list[str]:
    """Re-score documents with a cross-encoder and return top-k."""
    pairs = [[query, doc] for doc in documents]
    scores = reranker.predict(pairs)

    # Sort by score descending
    ranked = sorted(
        zip(documents, scores),
        key=lambda x: x[1],
        reverse=True,
    )
    return [doc for doc, score in ranked[:top_k]]

# Usage in the RAG pipeline:
# 1. Retrieve top-20 from ChromaDB
# 2. Rerank to top-5
# 3. Send top-5 to the LLM
initial_results = collection.query(query_embeddings=[q_emb], n_results=20)
reranked = rerank(question, initial_results["documents"][0], top_k=5)