from pathlib import Path

import torch
from sentence_transformers import SentenceTransformer


BASE_MODEL_NAME = "BAAI/bge-m3"
CHECKPOINT_FILENAME = "bge_m3_quantized_state_dict.pt"
MAX_SEQ_LENGTH = 128


def _quantize_model(model):
    try:
        from torch.ao.quantization import quantize_dynamic
    except ImportError:
        from torch.quantization import quantize_dynamic

    result = quantize_dynamic(
        model,
        {torch.nn.Linear},
        dtype=torch.qint8,
        inplace=True,
    )

    return result if result is not None else model


def load_model(checkpoint_path=None):
    if checkpoint_path is None:
        checkpoint_path = (
            Path(__file__).resolve().parent
            / CHECKPOINT_FILENAME
        )
    else:
        checkpoint_path = Path(checkpoint_path)

    if not checkpoint_path.is_file():
        raise FileNotFoundError(
            f"Checkpoint не найден: {checkpoint_path}"
        )

    model = SentenceTransformer(
        BASE_MODEL_NAME,
        device="cpu",
    )

    model.max_seq_length = MAX_SEQ_LENGTH
    model.eval()

    model = _quantize_model(model)
    model.max_seq_length = MAX_SEQ_LENGTH
    model.eval()

    try:
        state_dict = torch.load(
            checkpoint_path,
            map_location="cpu",
            weights_only=False,
        )
    except TypeError:
        state_dict = torch.load(
            checkpoint_path,
            map_location="cpu",
        )

    model.load_state_dict(
        state_dict,
        strict=True,
    )

    model.eval()
    return model


if __name__ == "__main__":
    model = load_model()

    embeddings = model.encode(
        ["Первый текст", "Второй текст"],
        batch_size=2,
        show_progress_bar=False,
        normalize_embeddings=True,
        convert_to_numpy=True,
        device="cpu",
    )

    print("Модель успешно загружена.")
    print("Размер эмбеддингов:", embeddings.shape)