model_loader.py

"""
IBM Granite 4.1 3B Instruct model loader and generation utilities.
"""

from __future__ import annotations

import os
from typing import Optional

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, Pipeline

from config import cfg
from logging_config import get_logger, setup_logging
from utils import Timer

setup_logging(log_dir=cfg.app.log_dir)
logger = get_logger(__name__)


class GraniteModelLoader:
    """
    Loads IBM Granite 4.1 3B Instruct and exposes a text-generation pipeline.
    """

    def __init__(self) -> None:
        self.pipe: Optional[Pipeline] = None
        self.tokenizer = None
        self.model_id = cfg.model.model_id
        self._is_loaded = False

    # ── Loading ───────────────────────────────────────────────────────────────

    def load(self, token: Optional[str] = None) -> None:
        if self._is_loaded:
            logger.info("Model already loaded; skipping.")
            return

        hf_token = token or cfg.hf_token
        if not hf_token:
            raise EnvironmentError(
                "Hugging Face token is required to load gated models. "
                "Set the HF_TOKEN environment variable."
            )

        logger.info("Loading tokenizer for '%s'…", self.model_id)
        with Timer() as t:
            self.tokenizer = AutoTokenizer.from_pretrained(
                self.model_id,
                token=hf_token,
            )
        logger.info("Tokenizer loaded in %s.", t)

        dtype = self._resolve_dtype()
        logger.info("Loading model with dtype=%s…", dtype)
        with Timer() as t:
            model = AutoModelForCausalLM.from_pretrained(
                self.model_id,
                token=hf_token,
                torch_dtype=dtype,
                device_map=cfg.model.device_map,
                low_cpu_mem_usage=True,
            )
        logger.info("Model loaded in %s.", t)

        self.pipe = pipeline(
            "text-generation",
            model=model,
            tokenizer=self.tokenizer,
            return_full_text=False,
        )
        self._is_loaded = True
        logger.info("Generation pipeline ready.")

    # ── Generation ────────────────────────────────────────────────────────────

    def generate(
        self,
        prompt: str,
        max_new_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        repetition_penalty: Optional[float] = None,
    ) -> str:
        if not self._is_loaded or self.pipe is None:
            raise RuntimeError("Model is not loaded. Call load() first.")

        params = {
            "max_new_tokens":    max_new_tokens    or cfg.model.max_new_tokens,
            "temperature":       temperature       or cfg.model.temperature,
            "top_p":             top_p             or cfg.model.top_p,
            "repetition_penalty": repetition_penalty or cfg.model.repetition_penalty,
            "do_sample":         cfg.model.do_sample,
        }

        logger.info(
            "Generating response (max_new_tokens=%d, temperature=%.2f)…",
            params["max_new_tokens"],
            params["temperature"],
        )

        with Timer() as t:
            output = self.pipe(prompt, **params)

        text = output[0]["generated_text"].strip()
        logger.info("Response generated in %s (%d chars).", t, len(text))
        return text

    # ── Prompt construction ───────────────────────────────────────────────────

    @staticmethod
    def build_prompt(
        query: str,
        retrieved_context: str,
        conversation_history: str = "",
    ) -> str:
        """
        Builds a retrieval-first prompt for IBM Granite chat models.
        Uses the model's expected <|...|> chat tokens.
        """
        system_msg = (
            "You are a professional university admissions assistant. "
            "Your role is to help prospective students with accurate information "
            "about admissions requirements, programs, fees, scholarships, deadlines, "
            "and related topics.\n\n"
            "STRICT RULES:\n"
            "1. Answer ONLY using the information provided in the CONTEXT section below.\n"
            "2. If the answer is not present in the context, respond: "
            "'I'm sorry, I couldn't find that information in the university knowledge base. "
            "Please contact the admissions office directly for assistance.'\n"
            "3. NEVER invent, guess, or extrapolate admissions policies, fees, or deadlines.\n"
            "4. Cite the source document name when possible.\n"
            "5. Keep answers concise, accurate, and professionally toned.\n"
            "6. If the query is off-topic (not related to university admissions), "
            "politely redirect the user."
        )

        context_block = (
            f"CONTEXT (retrieved from university knowledge base):\n"
            f"{'=' * 60}\n"
            f"{retrieved_context}\n"
            f"{'=' * 60}"
        )

        history_block = ""
        if conversation_history:
            history_block = (
                f"\nCONVERSATION HISTORY:\n{conversation_history}\n"
            )

        user_content = (
            f"{context_block}\n"
            f"{history_block}\n"
            f"QUESTION: {query}"
        )

        # IBM Granite 4.1 uses the standard HuggingFace chat template;
        # we format manually for pipeline compatibility.
        prompt = (
            f"<|start_of_role|>system<|end_of_role|>{system_msg}<|end_of_text|>\n"
            f"<|start_of_role|>user<|end_of_role|>{user_content}<|end_of_text|>\n"
            f"<|start_of_role|>assistant<|end_of_role|>"
        )
        return prompt

    # ── Helpers ───────────────────────────────────────────────────────────────

    @staticmethod
    def _resolve_dtype() -> torch.dtype:
        if torch.cuda.is_available():
            return torch.bfloat16
        if torch.backends.mps.is_available():
            return torch.float16
        return torch.float32

    @property
    def is_loaded(self) -> bool:
        return self._is_loaded

    @property
    def model_name(self) -> str:
        return self.model_id