04_vector_representation.py

"""
04_vector_representation.py
----------------------------
Stage 4 of the RAG pipeline: VECTOR REPRESENTATION (EMBEDDINGS).

Converts text chunks into numerical vector embeddings using a local
sentence-transformers model. A local embedding model is used so the
project does not require an embeddings-specific paid API key -- only the
OpenRouter key (used later, for generation) is needed.

Model: "all-MiniLM-L6-v2" (fast, small, good quality for student projects).

Run directly to see a demo of this stage:
    python 04_vector_representation.py
"""

from __future__ import annotations

import importlib.util
import sys
import os
from typing import List

_THIS_DIR = os.path.dirname(os.path.abspath(__file__))

EMBEDDING_MODEL_NAME = "all-MiniLM-L6-v2"

_model = None  # lazily loaded singleton


def _import_module_from_file(module_name: str, filename: str):
    path = os.path.join(_THIS_DIR, filename)
    spec = importlib.util.spec_from_file_location(module_name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module  # required so dataclasses can resolve cls.__module__
    spec.loader.exec_module(module)
    return module


_chunking_module = _import_module_from_file("chunking_stage", "03_chunking.py")
_preprocessing_module = _import_module_from_file("preprocessing_stage", "02_preprocessing.py")
_documents_module = _import_module_from_file("documents_stage", "01_documents.py")
Chunk = _chunking_module.Chunk


def get_embedding_model():
    """Load (once) and return the sentence-transformers embedding model."""
    global _model
    if _model is None:
        from sentence_transformers import SentenceTransformer

        _model = SentenceTransformer(EMBEDDING_MODEL_NAME)
    return _model


def embed_texts(texts: List[str]) -> List[List[float]]:
    """Compute embedding vectors for a list of raw strings."""
    if not texts:
        return []
    model = get_embedding_model()
    embeddings = model.encode(texts, show_progress_bar=False, convert_to_numpy=True)
    return embeddings.tolist()


def embed_chunks(chunks: List[Chunk]) -> List[List[float]]:
    """Compute embedding vectors for a list of Chunk objects."""
    texts = [chunk.text for chunk in chunks]
    return embed_texts(texts)


if __name__ == "__main__":
    raw_docs = _documents_module.load_documents()
    cleaned_docs = _preprocessing_module.preprocess_documents(raw_docs)
    chunks = _chunking_module.chunk_documents(cleaned_docs)

    if not chunks:
        print("No chunks available to embed. Add documents to the data/ folder first.")
    else:
        vectors = embed_chunks(chunks)
        print(f"Computed embeddings for {len(vectors)} chunk(s).")
        print(f"Embedding dimensionality: {len(vectors[0])}")
        print(f"First vector (truncated): {vectors[0][:8]}...")