70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
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 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()
|