scripts/embed.py

#!/usr/bin/env python3
"""Embedding generation CLI script for the RAG pipeline.

This script generates embeddings from text chunks and builds search indexes
for the pythermalcomfort RAG chatbot. It supports:
    - Reading chunks from JSONL files
    - Generating embeddings using BGE encoder with GPU acceleration
    - Building FAISS indexes for dense retrieval
    - Building BM25 indexes for sparse retrieval
    - Publishing artifacts to HuggingFace (optional)

Usage:
    # Basic embedding generation
    poetry run python scripts/embed.py data/chunks/chunks.jsonl data/embeddings/

    # With HuggingFace publishing
    poetry run python scripts/embed.py data/chunks/chunks.jsonl data/embeddings/ \
        --publish

    # Custom batch size and model
    poetry run python scripts/embed.py data/chunks/chunks.jsonl data/embeddings/ \
        --batch-size 64 --model BAAI/bge-base-en-v1.5

Output Files:
    {output_dir}/
    ├── embeddings.parquet      # Embeddings with chunk_id mapping
    ├── metadata.json           # Model metadata
    ├── faiss_index.bin         # FAISS index for dense retrieval
    ├── faiss_index.bin.ids.json # Chunk ID mapping for FAISS
    ├── bm25_index.pkl          # BM25 index for sparse retrieval
    └── chunks.parquet          # Chunks in parquet format (if --publish)

"""

from __future__ import annotations

import argparse
import json
import logging
import sys
import time
from pathlib import Path
from typing import TYPE_CHECKING

# =============================================================================
# Environment Variable Loading
# =============================================================================
# Load environment variables from .env file at script startup.
# This ensures HF_TOKEN and other secrets are available before they're needed.
# The .env file should be in the project root directory.
# =============================================================================
from dotenv import load_dotenv

# Find the project root (parent of scripts/ directory) and load .env from there
_PROJECT_ROOT = Path(__file__).parent.parent
_ENV_FILE = _PROJECT_ROOT / ".env"

if _ENV_FILE.exists():
    load_dotenv(_ENV_FILE)

# Rich is lightweight - import at module level for progress display
from rich.console import Console
from rich.progress import (
    BarColumn,
    MofNCompleteColumn,
    Progress,
    SpinnerColumn,
    TaskID,
    TaskProgressColumn,
    TextColumn,
    TimeElapsedColumn,
    TimeRemainingColumn,
)
from rich.table import Table

if TYPE_CHECKING:
    from rag_chatbot.chunking.models import Chunk
    from rag_chatbot.embeddings import (
        BGEEncoder,
        EmbeddingRecord,
    )

# =============================================================================
# Module Constants
# =============================================================================

# Default embedding model (BAAI General Embedding - small variant)
DEFAULT_MODEL: str = "BAAI/bge-small-en-v1.5"

# Default batch size for embedding generation
DEFAULT_BATCH_SIZE: int = 32

# Default embedding dimension for bge-small-en-v1.5
BGE_SMALL_DIM: int = 384

# =============================================================================
# Logging Configuration
# =============================================================================

# Configure logging to stderr with timestamp
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler(sys.stderr)],
)
logger = logging.getLogger(__name__)

# Rich console for styled output
console = Console()


# =============================================================================
# Helper Functions
# =============================================================================


def get_device_info() -> tuple[str, str]:
    """Get GPU/device information for reporting.

    This function checks for CUDA availability and returns device information
    for display in the CLI output. It imports torch lazily to avoid loading
    heavy dependencies until needed.

    Returns:
    -------
        Tuple of (device_type, device_name) where:
        - device_type: "cuda" or "cpu"
        - device_name: GPU name (e.g., "NVIDIA RTX 4090") or "CPU"

    Example:
    -------
        >>> device_type, device_name = get_device_info()
        >>> print(f"Using {device_name}")
        Using NVIDIA RTX 4090

    """
    # Lazy import torch to avoid loading until needed
    import torch  # type: ignore[import-not-found]

    if torch.cuda.is_available():
        device_type = "cuda"
        device_name = torch.cuda.get_device_name(0)
    else:
        device_type = "cpu"
        device_name = "CPU"

    return device_type, device_name


def load_chunks_from_jsonl(input_path: Path) -> list[Chunk]:
    """Load chunks from a JSONL file.

    Reads a JSONL file containing chunk data and parses each line into
    a Chunk model instance. Text normalization is applied during loading
    to fix common PDF extraction artifacts.

    Args:
    ----
        input_path: Path to the chunks.jsonl file.

    Returns:
    -------
        List of Chunk objects parsed from the file.

    Raises:
    ------
        FileNotFoundError: If the input file doesn't exist.
        ValueError: If the file is empty or contains invalid JSON.

    Example:
    -------
        >>> chunks = load_chunks_from_jsonl(Path("data/chunks/chunks.jsonl"))
        >>> len(chunks)
        1500

    """
    # Lazy import to avoid loading heavy dependencies at module level
    from rag_chatbot.chunking.models import Chunk, TextNormalizer

    if not input_path.exists():
        msg = f"Input file not found: {input_path}"
        raise FileNotFoundError(msg)

    chunks: list[Chunk] = []
    normalizer = TextNormalizer()

    with open(input_path, encoding="utf-8") as f:
        for line_num, raw_line in enumerate(f, start=1):
            # Skip empty lines
            line = raw_line.strip()
            if not line:
                continue

            try:
                # Parse JSON line
                data = json.loads(line)

                # Normalize text before creating chunk
                # Apply text normalization to fix OCR artifacts
                if "text" in data:
                    data["text"] = normalizer.normalize(data["text"], is_heading=False)

                # Create Chunk instance
                chunk = Chunk(**data)
                chunks.append(chunk)

            except json.JSONDecodeError as exc:
                logger.warning("Invalid JSON at line %d: %s", line_num, exc)
                continue
            except Exception as exc:
                logger.warning("Error parsing chunk at line %d: %s", line_num, exc)
                continue

    if not chunks:
        msg = f"No valid chunks found in {input_path}"
        raise ValueError(msg)

    return chunks


def create_embedding_records(
    chunks: list[Chunk],
    encoder: BGEEncoder,
    batch_size: int,
    progress: Progress,
    task_id: TaskID,
) -> list[EmbeddingRecord]:
    """Generate embeddings for chunks and create EmbeddingRecord objects.

    This function encodes all chunk texts using the BGE encoder and creates
    EmbeddingRecord instances with chunk_id, chunk_hash, and embedding data.
    Progress is reported through the Rich progress bar.

    Args:
    ----
        chunks: List of Chunk objects to embed.
        encoder: BGEEncoder instance for generating embeddings.
        batch_size: Number of chunks to process per batch.
        progress: Rich Progress instance for progress tracking.
        task_id: Task ID for the progress bar.

    Returns:
    -------
        List of EmbeddingRecord objects with generated embeddings.

    Example:
    -------
        >>> records = create_embedding_records(chunks, encoder, 32, progress, task_id)
        >>> len(records) == len(chunks)
        True

    """
    # Lazy import
    from rag_chatbot.embeddings import EmbeddingRecord

    # Extract texts for embedding
    texts = [chunk.text for chunk in chunks]

    # Track progress through callback
    def progress_callback(current_batch: int, _total_batches: int) -> None:
        """Update progress bar after each batch."""
        progress.update(task_id, completed=current_batch)

    # Generate embeddings with progress tracking
    embeddings = encoder.encode(
        texts=texts,
        batch_size=batch_size,
        show_progress=False,  # We use our own progress bar
        progress_callback=progress_callback,
    )

    # Create EmbeddingRecord for each chunk
    records: list[EmbeddingRecord] = []
    for idx, chunk in enumerate(chunks):
        record = EmbeddingRecord(
            chunk_id=chunk.chunk_id,
            chunk_hash=chunk.chunk_hash,
            embedding=embeddings[idx].tolist(),
        )
        records.append(record)

    return records


def build_indexes(
    output_dir: Path,
    chunks: list[Chunk],
    progress: Progress,
) -> tuple[float, float]:
    """Build FAISS and BM25 indexes from embeddings and chunks.

    This function builds both dense (FAISS) and sparse (BM25) indexes
    for hybrid retrieval. The FAISS index is built from the embeddings
    parquet file, while BM25 is built from chunk texts.

    Args:
    ----
        output_dir: Directory containing embeddings.parquet and for saving indexes.
        chunks: List of Chunk objects for BM25 indexing.
        progress: Rich Progress instance for progress tracking.

    Returns:
    -------
        Tuple of (faiss_build_time, bm25_build_time) in seconds.

    Example:
    -------
        >>> faiss_time, bm25_time = build_indexes(output_dir, chunks, progress)
        >>> print(f"FAISS: {faiss_time:.2f}s, BM25: {bm25_time:.2f}s")

    """
    # Lazy imports
    from rag_chatbot.embeddings import BM25IndexBuilder, FAISSIndexBuilder

    embeddings_path = output_dir / "embeddings.parquet"

    # Build FAISS index
    faiss_task = progress.add_task("[cyan]Building FAISS index...", total=1)
    faiss_start = time.perf_counter()

    faiss_builder = FAISSIndexBuilder()
    faiss_index = faiss_builder.build_from_parquet(embeddings_path)
    faiss_builder.save_index(faiss_index, output_dir / "faiss_index.bin")

    faiss_time = time.perf_counter() - faiss_start
    progress.update(faiss_task, completed=1)

    # Build BM25 index
    bm25_task = progress.add_task("[cyan]Building BM25 index...", total=1)
    bm25_start = time.perf_counter()

    bm25_builder = BM25IndexBuilder()
    bm25_index, chunk_ids = bm25_builder.build_from_chunks(chunks)
    bm25_builder.save_index(bm25_index, chunk_ids, output_dir / "bm25_index.pkl")

    bm25_time = time.perf_counter() - bm25_start
    progress.update(bm25_task, completed=1)

    return faiss_time, bm25_time


def publish_to_huggingface(
    output_dir: Path,
    chunks: list[Chunk],
    model_name: str,
    embedding_dim: int,
    progress: Progress,
) -> str:
    """Publish all artifacts to HuggingFace dataset repository.

    This function handles the complete publishing workflow:
    1. Saves chunks to parquet format
    2. Generates source manifest
    3. Authenticates with HuggingFace
    4. Uploads all artifacts

    Args:
    ----
        output_dir: Directory containing artifacts to publish.
        chunks: List of Chunk objects for chunks.parquet.
        model_name: Name of the embedding model used.
        embedding_dim: Dimension of embeddings.
        progress: Rich Progress instance for progress tracking.

    Returns:
    -------
        URL of the published HuggingFace dataset.

    Raises:
    ------
        ValueError: If HF_TOKEN is not set.
        RuntimeError: If publishing fails.

    Example:
    -------
        >>> url = publish_to_huggingface(
        ...     output_dir, chunks, "BAAI/bge-small-en-v1.5", 384, progress
        ... )
        >>> print(url)
        'https://huggingface.co/datasets/sadickam/pytherm_index'

    """
    # Lazy imports
    from rag_chatbot.embeddings import HuggingFacePublisher, PublisherConfig

    publish_task = progress.add_task("[cyan]Publishing to HuggingFace...", total=4)

    # Step 1: Save chunks to parquet
    config = PublisherConfig()
    publisher = HuggingFacePublisher(config)

    publisher.save_chunks_parquet(chunks, output_dir)
    progress.update(publish_task, advance=1)

    # Step 2: Generate source manifest
    manifest = publisher.generate_source_manifest(
        source_files=[],  # Source files not tracked in this context
        total_chunks=len(chunks),
        total_embeddings=len(chunks),
    )
    progress.update(publish_task, advance=1)

    # Step 3: Authenticate
    publisher.authenticate()
    progress.update(publish_task, advance=1)

    # Step 4: Publish all artifacts
    dataset_url = publisher.publish(
        artifacts_dir=output_dir,
        manifest=manifest,
        model_name=model_name,
        embedding_dimension=embedding_dim,
    )
    progress.update(publish_task, advance=1)

    return dataset_url


def print_statistics(  # noqa: PLR0913
    total_chunks: int,
    total_time: float,
    embedding_time: float,
    faiss_time: float,
    bm25_time: float,
    device_name: str,
    model_name: str,
    output_dir: Path,
    dataset_url: str | None = None,
) -> None:
    """Print final statistics table using Rich.

    Displays a formatted table with embedding statistics including:
    - Total chunks processed
    - Time breakdowns (embedding, indexing)
    - Throughput metrics
    - Device information
    - Output file sizes

    Args:
    ----
        total_chunks: Number of chunks embedded.
        total_time: Total elapsed time in seconds.
        embedding_time: Time spent on embedding generation.
        faiss_time: Time spent building FAISS index.
        bm25_time: Time spent building BM25 index.
        device_name: Name of the device used (GPU name or "CPU").
        model_name: Name of the embedding model.
        output_dir: Directory where outputs were saved.
        dataset_url: Optional URL of published HuggingFace dataset.

    """
    # Calculate throughput
    throughput = total_chunks / embedding_time if embedding_time > 0 else 0

    # Create statistics table
    table = Table(
        title="Embedding Statistics", show_header=True, header_style="bold cyan"
    )
    table.add_column("Metric", style="dim", width=25)
    table.add_column("Value", justify="right")

    # Add rows
    table.add_row("Total Chunks", f"{total_chunks:,}")
    table.add_row("Model", model_name)
    table.add_row("Device", device_name)
    table.add_row("", "")  # Separator
    table.add_row("Embedding Time", f"{embedding_time:.2f}s")
    table.add_row("FAISS Build Time", f"{faiss_time:.2f}s")
    table.add_row("BM25 Build Time", f"{bm25_time:.2f}s")
    table.add_row("Total Time", f"{total_time:.2f}s")
    table.add_row("", "")  # Separator
    table.add_row("Throughput", f"{throughput:.1f} chunks/sec")

    # Add file sizes
    embeddings_file = output_dir / "embeddings.parquet"
    faiss_file = output_dir / "faiss_index.bin"
    bm25_file = output_dir / "bm25_index.pkl"

    if embeddings_file.exists():
        size_mb = embeddings_file.stat().st_size / (1024 * 1024)
        table.add_row("Embeddings Size", f"{size_mb:.2f} MB")

    if faiss_file.exists():
        size_mb = faiss_file.stat().st_size / (1024 * 1024)
        table.add_row("FAISS Index Size", f"{size_mb:.2f} MB")

    if bm25_file.exists():
        size_mb = bm25_file.stat().st_size / (1024 * 1024)
        table.add_row("BM25 Index Size", f"{size_mb:.2f} MB")

    # Add dataset URL if published
    if dataset_url:
        table.add_row("", "")  # Separator
        table.add_row("Published URL", dataset_url)

    console.print()
    console.print(table)
    console.print()


def parse_args() -> argparse.Namespace:
    """Parse command line arguments.

    Returns
    -------
        Parsed argument namespace with input_path, output_dir, and options.

    """
    parser = argparse.ArgumentParser(
        description="Generate embeddings and build indexes for RAG retrieval.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  poetry run python scripts/embed.py chunks.jsonl output/
  poetry run python scripts/embed.py chunks.jsonl output/ --publish
  poetry run python scripts/embed.py chunks.jsonl output/ --batch-size 64
        """,
    )

    parser.add_argument(
        "input_path",
        type=Path,
        help="Path to chunks.jsonl file containing text chunks",
    )

    parser.add_argument(
        "output_dir",
        type=Path,
        help="Directory to save embeddings and indexes",
    )

    parser.add_argument(
        "--publish",
        action="store_true",
        help="Publish artifacts to HuggingFace after embedding",
    )

    parser.add_argument(
        "--batch-size",
        type=int,
        default=DEFAULT_BATCH_SIZE,
        help=f"Batch size for embedding generation (default: {DEFAULT_BATCH_SIZE})",
    )

    parser.add_argument(
        "--model",
        type=str,
        default=DEFAULT_MODEL,
        help=f"Embedding model name (default: {DEFAULT_MODEL})",
    )

    parser.add_argument(
        "--device",
        type=str,
        default="auto",
        choices=["cpu", "cuda", "auto"],
        help="Device to use for embedding (default: auto)",
    )

    return parser.parse_args()


def main() -> int:
    """Run the embedding generation pipeline.

    This is the main entry point for the embed.py CLI script. It orchestrates:
    1. Loading chunks from JSONL
    2. Initializing the encoder with GPU if available
    3. Generating embeddings with progress tracking
    4. Saving embeddings to parquet storage
    5. Building FAISS and BM25 indexes
    6. Optionally publishing to HuggingFace

    Returns
    -------
        Exit code (0 for success, 1 for error).

    """
    # Parse command line arguments
    args = parse_args()

    # Track total time
    start_time = time.perf_counter()

    try:
        # =================================================================
        # Step 1: Display GPU/Device Information
        # =================================================================
        console.print("\n[bold cyan]Embedding Generation Pipeline[/bold cyan]\n")

        device_type, device_name = get_device_info()
        device = args.device if args.device != "auto" else None

        console.print(f"[green]Device:[/green] {device_name}")
        console.print(f"[green]Model:[/green] {args.model}")
        console.print(f"[green]Batch Size:[/green] {args.batch_size}")
        console.print()

        # =================================================================
        # Step 2: Load Chunks from JSONL
        # =================================================================
        console.print(f"[cyan]Loading chunks from {args.input_path}...[/cyan]")

        chunks = load_chunks_from_jsonl(args.input_path)
        console.print(f"[green]Loaded {len(chunks):,} chunks[/green]\n")

        # Handle empty input gracefully
        if not chunks:
            console.print("[yellow]Warning: No chunks to process. Exiting.[/yellow]")
            return 0

        # =================================================================
        # Step 3: Initialize Encoder and Storage
        # =================================================================
        # Lazy imports for heavy dependencies
        from rag_chatbot.embeddings import (
            BGEEncoder,
            EmbeddingBatch,
            EmbeddingStorage,
        )

        # Create output directory
        args.output_dir.mkdir(parents=True, exist_ok=True)

        # Initialize encoder with configured model and device
        encoder = BGEEncoder(
            model_name=args.model,
            device=device,
            normalize_text=False,  # Already normalized during chunk loading
        )

        # Initialize storage
        storage = EmbeddingStorage(args.output_dir)

        # =================================================================
        # Step 4: Generate Embeddings with Progress Tracking
        # =================================================================
        # Calculate total batches for progress bar
        import math

        total_batches = math.ceil(len(chunks) / args.batch_size)

        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TaskProgressColumn(),
            MofNCompleteColumn(),
            TimeElapsedColumn(),
            TimeRemainingColumn(),
            console=console,
        ) as progress:
            # Embedding task
            embed_task = progress.add_task(
                "[cyan]Embedding chunks...",
                total=total_batches,
            )

            embedding_start = time.perf_counter()

            # Generate embeddings
            records = create_embedding_records(
                chunks=chunks,
                encoder=encoder,
                batch_size=args.batch_size,
                progress=progress,
                task_id=embed_task,
            )

            embedding_time = time.perf_counter() - embedding_start

            # Complete the embedding task
            progress.update(embed_task, completed=total_batches)

            # =================================================================
            # Step 5: Save Embeddings to Storage
            # =================================================================
            save_task = progress.add_task("[cyan]Saving embeddings...", total=1)

            # Create EmbeddingBatch
            batch = EmbeddingBatch(
                model_name=args.model,
                dimension=encoder.embedding_dim,
                dtype="float16",
                records=records,
            )

            # Save to storage
            storage.save(batch)

            progress.update(save_task, completed=1)

            # =================================================================
            # Step 6: Build FAISS and BM25 Indexes
            # =================================================================
            faiss_time, bm25_time = build_indexes(
                output_dir=args.output_dir,
                chunks=chunks,
                progress=progress,
            )

            # =================================================================
            # Step 7: Publish to HuggingFace (if requested)
            # =================================================================
            dataset_url: str | None = None

            if args.publish:
                dataset_url = publish_to_huggingface(
                    output_dir=args.output_dir,
                    chunks=chunks,
                    model_name=args.model,
                    embedding_dim=encoder.embedding_dim,
                    progress=progress,
                )

        # =================================================================
        # Step 8: Print Statistics
        # =================================================================
        total_time = time.perf_counter() - start_time

        print_statistics(
            total_chunks=len(chunks),
            total_time=total_time,
            embedding_time=embedding_time,
            faiss_time=faiss_time,
            bm25_time=bm25_time,
            device_name=device_name,
            model_name=args.model,
            output_dir=args.output_dir,
            dataset_url=dataset_url,
        )

    except FileNotFoundError as exc:
        console.print(f"[bold red]Error:[/bold red] {exc}")
        return 1
    except ValueError as exc:
        console.print(f"[bold red]Error:[/bold red] {exc}")
        return 1
    except KeyboardInterrupt:
        console.print("\n[yellow]Interrupted by user[/yellow]")
        return 1
    except Exception as exc:
        console.print(f"[bold red]Unexpected error:[/bold red] {exc}")
        logger.exception("Unexpected error during embedding generation")
        return 1
    else:
        console.print("[bold green]Embedding generation complete![/bold green]\n")
        return 0


if __name__ == "__main__":
    sys.exit(main())