AIServ/models/model_loader.py

import torch
from sentence_transformers import SentenceTransformer
import joblib
import os

class ModelLoader:
    def __init__(self):
        self.text_model = None
        self.classification_model = None
        self.summarization_model = None
        self.summarization_tokenizer = None
        
    def load_text_model(self, use_local=True, local_model_path="./models/bin/all-MiniLM-L6-v2"):
        if self.text_model is None:
            if use_local and os.path.exists(local_model_path):
                print("🔧 从本地加载sentence-transformers模型...ok")
                self.text_model = SentenceTransformer(local_model_path)
            else:
                print("🔧 从网络加载sentence-transformers模型...ok")
                self.text_model = SentenceTransformer('all-MiniLM-L6-v2')
        return self.text_model

# 全局模型加载器
model_loader = ModelLoader()