Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import torch | |
| from transformers import AutoTokenizer, BertTokenizer | |
| from huggingface_hub import snapshot_download | |
| from models.classifier import Classifier | |
| from models.sdg_classifier import BERTClassifier | |
| from configs.sdg_labels import SDG_LABELS | |
| from configs.model_config import MODEL_FOLDERS | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| # MODEL PATHS | |
| BASE_PATH = "./models/" | |
| # LOAD MODEL FROM FOLDER (original models) | |
| def load_model_from_folder(folder_path): | |
| # CONFIG | |
| with open(os.path.join(folder_path, "config.json")) as f: | |
| config = json.load(f) | |
| model_name = config["model_name"] | |
| dropout = config.get("dropout_rate", 0.1) | |
| # LABELS (always normal classification) | |
| with open(os.path.join(folder_path, "labels.json")) as f: | |
| label_data = json.load(f) | |
| labels = label_data["label_list"] | |
| num_labels = len(labels) | |
| # TOKENIZER | |
| tokenizer = AutoTokenizer.from_pretrained(folder_path) | |
| # MODEL | |
| model = Classifier(model_name, num_labels, dropout) | |
| # WEIGHTS | |
| state = torch.load(os.path.join(folder_path, "model.pt"), map_location=device) | |
| model.load_state_dict(state, strict=True) | |
| model.to(device) | |
| model.eval() | |
| return { | |
| "model": model, | |
| "tokenizer": tokenizer, | |
| "labels": labels, | |
| "config": config | |
| } | |
| def load_sdg_model_from_folder(folder_path): | |
| # CONFIG | |
| with open(os.path.join(folder_path, "config.json")) as f: | |
| config = json.load(f) | |
| dropout = config.get("dropout_rate", 0.1) | |
| num_classes = config["num_classes"] # 17 | |
| # LABEL MAP {"0": "1", "1": "2", ..., "16": "17"} | |
| with open(os.path.join(folder_path, "label_map.json")) as f: | |
| label_map = json.load(f) | |
| # Convert to ordered list by integer key -> ["1", "2", ..., "17"] | |
| raw_labels = [label_map[str(i)] for i in range(num_classes)] | |
| pretty_labels = [SDG_LABELS[x] for x in raw_labels] | |
| # TOKENIZER (saved inside a 'tokenizer' subfolder) | |
| tokenizer_path = os.path.join(folder_path, "tokenizer") | |
| tokenizer = BertTokenizer.from_pretrained(tokenizer_path) | |
| # MODEL | |
| model = BERTClassifier(n_classes=num_classes, dropout_rate=dropout) | |
| # WEIGHTS | |
| state = torch.load(os.path.join(folder_path, "model_state.pt"), map_location=device) | |
| model.load_state_dict(state, strict=True) | |
| model.to(device) | |
| model.eval() | |
| return { | |
| "model": model, | |
| "tokenizer": tokenizer, | |
| "labels": raw_labels, # for prediction logic | |
| "pretty_labels": pretty_labels, # for charts/UI | |
| "config": config, | |
| "is_sdg": True # flag so prediction knows which path to take | |
| } | |
| # MODEL CACHE | |
| LOADED_MODELS = {} | |
| def get_model(task_name): | |
| if task_name in LOADED_MODELS: | |
| return LOADED_MODELS[task_name] | |
| repo = f"sag-uniroma2/{MODEL_FOLDERS[task_name]}" | |
| # sag-uniroma2 models are public, so a token isn't required. | |
| # If HF_TOKEN is set (e.g. for higher rate limits or future-proofing | |
| # against the repo becoming gated/private later), it will still be used. | |
| token = os.getenv("HF_TOKEN") # may be None -- fine for public repos | |
| folder = snapshot_download(repo_id=repo, token=token) | |
| if task_name == "17 SDG Alignment": | |
| bundle = load_sdg_model_from_folder(folder) | |
| else: | |
| bundle = load_model_from_folder(folder) | |
| LOADED_MODELS[task_name] = bundle | |
| return bundle | |