model_setup.py

"""
Model loading for HuggingFace Spaces ZeroGPU deployment.
"""
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig,
    pipeline,
)
from langchain_huggingface import HuggingFaceEmbeddings

from config import MODEL_ID, EMBEDDING_MODEL_ID, MAX_NEW_TOKENS
from logging_config import get_logger

logger = get_logger(__name__)

# --- Tokenizer -------------------------------
logger.info(f"Loading tokenizer: {MODEL_ID}")
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    trust_remote_code=True, # Qwen has a custom tokenizer class
    clean_up_tokenization_spaces=False, # BPE tokenizers must NOT strip spaces before punctuation — it
                                        # corrupts output. This flag is meant for WordPiece (BERT) only.
)

# --- Quantization config -----------------------------------
# 4-bit NF4 (Normal Float 4-bit) is the most accurate 4-bit format for LLM weights.
# double_quant saves an additional ~0.4 bits per parameter on the quantization constants themselves.
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

# --- Model --------------------------------------------------
# device_map={"": 0} pins the ENTIRE model onto GPU slot 0 at load time.
# This is the module-scope CUDA placement ZeroGPU requires — the
# `spaces` package (imported first in app.py) intercepts this and
# handles it correctly even though no physical GPU is attached yet
# during this import.
logger.info(f"Loading model (4-bit NF4 quantized) — device_map={{'':0}}…")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=quantization_config,
    device_map={"": 0}, # Module-scope GPU placement (see docstring)
    low_cpu_mem_usage=True, # Streams weights from disk, avoids RAM spike
    trust_remote_code=True,
)
model.eval() # Disable dropout — we're doing inference only

# NOTE: No torch.compile() here.
# ZeroGPU does not support torch.compile. Calling it would either
# silently fail to provide speedup or raise an error when the first
# @spaces.GPU-decorated call tries to actually run on a real GPU worker.

# --- Text generation pipeline --------------------------------
# The pipeline object itself does not run any CUDA kernels at
# construction time — it only wraps tokenizer + model + generation
# config. Building it here at module scope is safe. The ACTUAL
# inference call (invoking this pipeline) happens inside the
# @spaces.GPU-decorated function in generation.py — never here.
logger.info("Building text generation pipeline…")
text_gen_pipeline = pipeline(
    task="text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=MAX_NEW_TOKENS,
    max_length=None, # Override model's default max_length (often 20) to silence the HF conflict warning
    do_sample=False, # Greedy decoding: deterministic, faithful to context
    return_full_text=False, # Return only newly generated tokens, not the prompt
)

# --- Embedding model ---------------------------------------------
# MiniLM stays on CPU — it's fast enough there and doesn't need the
# GPU slot, which should be reserved for the LLM's generation calls.
logger.info("Loading embedding model…")
embedding_model = HuggingFaceEmbeddings(
    model_name=EMBEDDING_MODEL_ID,
    model_kwargs={"device": "cpu"},
    encode_kwargs={"normalize_embeddings": True},
)

logger.info("All models ready.")