labor-arbitration-system/nlp-service/app/services/model_loader.py

"""
Singleton model loader for GPU memory-efficient inference.
"""

from __future__ import annotations

import os
from typing import Optional

# NOTE: must be imported before torch on Windows
import transformers  # noqa: F401

import torch

from .bert_multi_task_model import ChineseRobertaMultiTask


class ModelLoader:
    """Thread-safe singleton for loading the BERT multi-task model once."""

    _instance: Optional["ModelLoader"] = None
    _model: Optional[ChineseRobertaMultiTask] = None
    _model_path: str = ""

    def __new__(cls) -> "ModelLoader":
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def get_model(self, model_path: str | None = None) -> ChineseRobertaMultiTask:
        path = model_path or os.environ.get(
            "MODEL_PATH",
            os.path.join(os.path.dirname(__file__), "..", "..", "models", "chinese_roberta_labor_extractor"),
        )

        if self._model is not None and path == self._model_path:
            return self._model

        self._model_path = path
        self._model = ChineseRobertaMultiTask.from_pretrained(path)
        self._model.eval()

        if torch.cuda.is_available():
            self._model.cuda()

        return self._model

    def clear(self):
        """Release model from memory."""
        self._model = None
        self._model_path = ""
        if torch.cuda.is_available():
            torch.cuda.empty_cache()