src/generation.py

"""
Generation engine for the Pleias-RAG model using llama.cpp backend.
Handles prompt formatting and raw token-by-token streaming.
"""

import logging
import time
from typing import Any, Dict, Iterator, List

logger = logging.getLogger(__name__)


class GenerationEngine:
    """
    Engine for generating responses using a local GGUF model via llama.cpp.
    Formats prompts with special tokens and streams raw generated text.
    """

    def __init__(
        self,
        model_path_or_name: str,
        max_tokens: int = 2048,
        temperature: float = 0.1,
        top_p: float = 0.95,
        repetition_penalty: float = 1.0,
    ):
        """
        Initialize the generation engine with model parameters.

        Args:
            model_path_or_name: Path to the GGUF model file.
            max_tokens: Maximum number of tokens to generate.
            temperature: Sampling temperature (lower = more deterministic).
            top_p: Nucleus sampling probability threshold.
            repetition_penalty: Penalty for repeating tokens.
        """
        self.model_path = model_path_or_name
        self.max_tokens = max_tokens
        self.temperature = temperature
        self.top_p = top_p
        self.repetition_penalty = repetition_penalty
        self._init_llama_cpp()

    def _init_llama_cpp(self):
        """
        Load the model using llama.cpp backend.
        Configures context size, GPU layers, and thread count.
        """
        from llama_cpp import Llama

        logger.info("Loading model with llama_cpp")
        self.model = Llama(
            model_path=self.model_path,
            n_ctx=4096,       # Context window size
            n_gpu_layers=0,   # CPU only (set > 0 for GPU acceleration)
            verbose=False,
            n_threads=4,
            n_batch=512,      # Batch size for prompt processing
            use_mmap=True,    # Memory-map model for faster loading
            use_mlock=False,  # Don't lock in RAM (Pi5 has limited memory)
        )
        logger.info("Model loaded successfully!!!")

    def format_prompt(self, query: str, sources: List[Dict[str, Any]]) -> str:
        """
        Format the query and sources into a prompt using the Pleias-RAG ChatML format.

        The prompt structure is:
            <|im_start|>user
            {query}

            **source_1**
            {source_text}

            **source_2**
            ...
            <|im_end|>
            <|im_start|>assistant
            <think>

        Args:
            query: The user's question.
            sources: List of source documents, each with a "text" key.

        Returns:
            Formatted prompt string ready for tokenization.
        """
        prompt = f"<|im_start|>user\n{query}\n\n"

        # Add each source with its ID in **source_N** format
        for idx, source in enumerate(sources, 1):
            source_text = source.get("text", "")
            prompt += f"**source_{idx}**\n{source_text}\n\n"

        # End user turn and start assistant turn with <think> tag
        prompt += "<|im_end|>\n<|im_start|>assistant\n<think>\n"

        logger.debug(f"Formatted prompt: \n {prompt}")
        return prompt

    def stream_generate(self, query: str, sources: List[Dict[str, Any]]) -> Iterator[str]:
        """
        Stream the model's raw output token-by-token.

        Tokenizes the prompt with special=True to preserve special tokens,
        then yields each detokenized piece until a stop condition is met:
            - <|end_of_text|> token
            - <|im_end|> token
            - max_tokens limit reached

        Args:
            query: The user's question.
            sources: List of source documents retrieved from the database.

        Yields:
            Raw text pieces of the model output, in generation order.
        """
        formatted_prompt = self.format_prompt(query, sources)

        t0 = time.time()
        logger.info("Starting streaming generation...")

        tokens = self.model.generate(
            self.model.tokenize(formatted_prompt.encode("utf-8"), special=True),
            temp=self.temperature,
            top_p=self.top_p,
            repeat_penalty=self.repetition_penalty,
            reset=True,
        )

        t1 = None
        for i, t in enumerate(tokens):
            # Log time to first token (prefill time)
            if t1 is None:
                t1 = time.time()
                logger.info(f"Prefill time (time to first token): {t1 - t0:.2f} seconds")

            # Detokenize with special=True to render special tokens in output
            piece = self.model.detokenize([t], special=True).decode("utf-8", errors="replace")

            # Stop conditions
            if piece in ("<|end_of_text|>", "<|im_end|>") or i >= self.max_tokens:
                break

            yield piece

        logger.info(f"Total streaming generation time: {time.time() - t0:.2f} seconds")