app/model_loader.py

"""
Model Loader - Tải model và tokenizer từ HuggingFace Hub.

Module này chịu trách nhiệm:
1. Tải model + tokenizer DUY NHẤT MỘT LẦN khi app khởi động
2. Chuyển model sang eval() mode để tối ưu inference
3. Cung cấp hàm generate_embeddings() để trích xuất vector nhúng
"""

import torch
from transformers import AutoModel, AutoTokenizer

from app.config import settings


# Biến toàn cục - lưu trữ model và tokenizer sau khi load
model = None
tokenizer = None


def load_model():
    """
    Tải model và tokenizer từ HuggingFace Hub vào bộ nhớ.

    Hàm này được gọi DUY NHẤT MỘT LẦN khi FastAPI app khởi động.
    Model được chuyển sang eval() mode và đưa lên device phù hợp (CPU/GPU).
    """
    global model, tokenizer

    print(f"[INFO] Loading model: {settings.MODEL_NAME}")
    print(f"[INFO] Device: {settings.DEVICE}")

    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(
        settings.MODEL_NAME,
        token=settings.HF_TOKEN,
    )

    # Load model
    model = AutoModel.from_pretrained(
        settings.MODEL_NAME,
        token=settings.HF_TOKEN,
    )

    # Chuyển sang eval mode (tắt dropout, batch norm ở inference mode)
    model.eval()

    # Đưa model lên device (cpu hoặc cuda)
    model.to(settings.DEVICE)

    print("[INFO] Model loaded successfully! Ready to serve.")


def generate_embeddings(texts: list[str]) -> list[list[float]]:
    """
    Trích xuất vector nhúng từ danh sách văn bản.

    Pipeline xử lý:
    1. Tokenize input texts
    2. Feed tokens qua model (không tính gradient)
    3. Mean pooling trên token embeddings (bỏ qua padding tokens)
    4. L2 normalize vector kết quả

    Args:
        texts: Danh sách các chuỗi văn bản cần embedding.

    Returns:
        Danh sách các vector nhúng (mỗi vector là list[float]).
    """
    # Bước 1: Tokenize
    encoded = tokenizer(
        texts,
        padding=True,
        truncation=True,
        max_length=settings.MAX_LENGTH,
        return_tensors="pt",
    )

    # Đưa input tensors lên cùng device với model
    encoded = {key: val.to(settings.DEVICE) for key, val in encoded.items()}

    # Bước 2: Forward pass (không tính gradient để tiết kiệm bộ nhớ)
    with torch.no_grad():
        outputs = model(**encoded)

    # Bước 3: Mean Pooling
    # outputs.last_hidden_state có shape: (batch_size, seq_length, hidden_dim)
    token_embeddings = outputs.last_hidden_state
    attention_mask = encoded["attention_mask"]

    # Mở rộng attention_mask để khớp shape với token_embeddings
    # attention_mask: (batch_size, seq_length) -> (batch_size, seq_length, hidden_dim)
    mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()

    # Tính trung bình có trọng số (chỉ tính các token thực, bỏ qua padding)
    sum_embeddings = torch.sum(token_embeddings * mask_expanded, dim=1)
    sum_mask = torch.clamp(mask_expanded.sum(dim=1), min=1e-9)
    mean_embeddings = sum_embeddings / sum_mask

    # Bước 4: L2 Normalize
    normalized = torch.nn.functional.normalize(mean_embeddings, p=2, dim=1)

    # Chuyển tensor về list Python
    return normalized.cpu().tolist()


def compute_cosine_similarity(vector1: list[float], vector2: list[float]) -> float:
    """
    Tinh cosine similarity giua 2 vector embedding.

    Cosine similarity do goc giua 2 vector trong khong gian nhieu chieu.
    Gia tri tra ve nam trong khoang [-1, 1]:
    - 1.0  = hoan toan giong nhau
    - 0.0  = khong lien quan
    - -1.0 = hoan toan doi lap

    Args:
        vector1: Vector embedding thu nhat (list[float]).
        vector2: Vector embedding thu hai (list[float]).

    Returns:
        Diem cosine similarity (float).
    """
    # Chuyen list Python thanh tensor
    v1 = torch.tensor(vector1, dtype=torch.float32)
    v2 = torch.tensor(vector2, dtype=torch.float32)

    # Tinh cosine similarity bang PyTorch
    similarity = torch.nn.functional.cosine_similarity(
        v1.unsqueeze(0), v2.unsqueeze(0)
    )

    return similarity.item()


# ==================== MINILM MODEL ====================

# Bien toan cuc MiniLM - tach biet hoan toan voi Gemma
minilm_model = None
minilm_tokenizer = None


def load_minilm():
    """
    Tai MiniLM model va tokenizer tu HuggingFace Hub.

    MiniLM chay tren CPU vi:
    - Model chi co 22M params, CPU van cuc nhanh (~5-10ms/request)
    - Tiet kiem GPU memory cho Gemma 300M (model chinh)
    """
    global minilm_model, minilm_tokenizer

    print(f"[INFO] Loading MiniLM: {settings.MODEL_NAME_MINILM}")
    print("[INFO] MiniLM device: cpu")

    # Load tokenizer
    minilm_tokenizer = AutoTokenizer.from_pretrained(
        settings.MODEL_NAME_MINILM,
    )

    # Load model
    minilm_model = AutoModel.from_pretrained(
        settings.MODEL_NAME_MINILM,
    )

    # Eval mode + CPU
    minilm_model.eval()
    minilm_model.to("cpu")

    print("[INFO] MiniLM loaded successfully!")


def generate_embeddings_minilm(texts: list[str]) -> list[list[float]]:
    """
    Trich xuat vector nhung bang MiniLM (384 chieu).

    Pipeline giong voi Gemma nhung:
    - Chay tren CPU (khong dung GPU)
    - Vector output co 384 chieu (thay vi 768 cua Gemma)
    - MAX_LENGTH ngan hon (256 thay vi 512)

    Args:
        texts: Danh sach van ban can embedding (search query, job title, ...).

    Returns:
        Danh sach vector nhung 384 chieu.
    """
    # Buoc 1: Tokenize
    encoded = minilm_tokenizer(
        texts,
        padding=True,
        truncation=True,
        max_length=settings.MAX_LENGTH_MINILM,
        return_tensors="pt",
    )

    # MiniLM chay tren CPU
    encoded = {key: val.to("cpu") for key, val in encoded.items()}

    # Buoc 2: Forward pass
    with torch.no_grad():
        outputs = minilm_model(**encoded)

    # Buoc 3: Mean Pooling
    token_embeddings = outputs.last_hidden_state
    attention_mask = encoded["attention_mask"]
    mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()

    sum_embeddings = torch.sum(token_embeddings * mask_expanded, dim=1)
    sum_mask = torch.clamp(mask_expanded.sum(dim=1), min=1e-9)
    mean_embeddings = sum_embeddings / sum_mask

    # Buoc 4: L2 Normalize
    normalized = torch.nn.functional.normalize(mean_embeddings, p=2, dim=1)

    return normalized.cpu().tolist()