| import os |
| import torch |
| from transformers import AutoModel, AutoTokenizer |
|
|
| class ModelLoader: |
| def __init__(self): |
| self.loaded_models = {} |
| self.available_models = { |
| 'vit-base': 'google/vit-base-patch16-224', |
| 'resnet-50': 'microsoft/resnet-50', |
| 'dinov2-base': 'facebook/dinov2-base', |
| 'flan-t5-large': 'google/flan-t5-large', |
| 'bert-base': 'bert-base-uncased' |
| } |
| |
| def load_model(self, model_name: str): |
| if model_name in self.loaded_models: |
| return self.loaded_models[model_name] |
| |
| model_path = self.available_models.get(model_name) |
| if not model_path: |
| return None |
| |
| model = AutoModel.from_pretrained(model_path) |
| self.loaded_models[model_name] = model |
| return model |
| |
| def list_models(self) -> list: |
| return list(self.available_models.keys()) |
| |
| def get_loaded_models(self) -> list: |
| return list(self.loaded_models.keys()) |
| |
| def unload_model(self, model_name: str): |
| if model_name in self.loaded_models: |
| del self.loaded_models[model_name] |
| return True |
| return False |
