76 lines
2.3 KiB
Python
76 lines
2.3 KiB
Python
import os
|
|
import torch
|
|
from transformers import AutoImageProcessor
|
|
from transformers import AutoModelForImageClassification
|
|
from transformers import AutoModelForSequenceClassification
|
|
from transformers import AutoTokenizer
|
|
|
|
from app.config.logging import logger
|
|
from app.config.settings import settings
|
|
|
|
|
|
class ModelManager:
|
|
|
|
def __init__(self):
|
|
self.device = None
|
|
self.text_model = None
|
|
self.text_tokenizer = None
|
|
self.image_model = None
|
|
self.image_processor = None
|
|
|
|
def _setup_hf_token(self):
|
|
if settings.HF_TOKEN:
|
|
os.environ["HF_TOKEN"] = settings.HF_TOKEN
|
|
os.environ["HUGGING_FACE_HUB_TOKEN"] = settings.HF_TOKEN
|
|
|
|
def load_device(self):
|
|
if settings.DEVICE == "auto":
|
|
self.device = torch.device(
|
|
"cuda" if torch.cuda.is_available() else "cpu"
|
|
)
|
|
else:
|
|
self.device = torch.device(settings.DEVICE)
|
|
|
|
def load_tokenizer(self):
|
|
logger.info("Loading tokenizer...")
|
|
self.text_tokenizer = AutoTokenizer.from_pretrained(
|
|
settings.TEXT_MODEL,
|
|
cache_dir=settings.MODEL_CACHE_DIR
|
|
)
|
|
logger.info("Tokenizer loaded.")
|
|
|
|
def load_model(self):
|
|
logger.info("Loading model...")
|
|
self.text_model = AutoModelForSequenceClassification.from_pretrained(
|
|
settings.TEXT_MODEL,
|
|
cache_dir=settings.MODEL_CACHE_DIR
|
|
)
|
|
self.text_model.to(self.device)
|
|
self.text_model.eval()
|
|
logger.debug("id2label: %s", self.text_model.config.id2label)
|
|
logger.info("Text model loaded.")
|
|
|
|
def load_image_model(self):
|
|
logger.info("Loading image model...")
|
|
self.image_processor = AutoImageProcessor.from_pretrained(
|
|
settings.IMAGE_MODEL,
|
|
cache_dir=settings.MODEL_CACHE_DIR
|
|
)
|
|
self.image_model = AutoModelForImageClassification.from_pretrained(
|
|
settings.IMAGE_MODEL,
|
|
cache_dir=settings.MODEL_CACHE_DIR
|
|
)
|
|
self.image_model.to(self.device)
|
|
self.image_model.eval()
|
|
logger.info("Image model loaded.")
|
|
|
|
def initialize(self):
|
|
self.load_device()
|
|
logger.info(f"Using device: {self.device}")
|
|
self.load_tokenizer()
|
|
self.load_model()
|
|
self.load_image_model()
|
|
|
|
|
|
model_manager = ModelManager()
|