Spaces:
Running
Running
| from pathlib import Path | |
| import pickle | |
| import torch | |
| import torch.nn as nn | |
| import pandas as pd | |
| import json | |
| from sentence_transformers import SentenceTransformer | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| HIDDEN_DIM = 768 | |
| DROPOUT = 0.1 | |
| class FlatEmbedMLP(nn.Module): | |
| def __init__(self, input_dim, n_classes, hidden_dim=HIDDEN_DIM, dropout=DROPOUT): | |
| super().__init__() | |
| self.net = nn.Sequential( | |
| nn.Linear(input_dim, hidden_dim), | |
| nn.ReLU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(hidden_dim, hidden_dim), | |
| nn.ReLU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(hidden_dim, n_classes), | |
| ) | |
| def forward(self, x): | |
| return self.net(x) | |
| _artifacts = None | |
| def _artifacts_root(): | |
| return Path(__file__).resolve().parents[1] / "training" / "artifacts" | |
| def load_artifacts(): | |
| global _artifacts | |
| if _artifacts is not None: | |
| return _artifacts | |
| artifacts_dir = _artifacts_root() | |
| with open(artifacts_dir / "label_maps" / "label_maps_embed.pkl", "rb") as f: | |
| label_maps = pickle.load(f) | |
| with open(artifacts_dir / "embedder" / "embed_metadata.pkl", "rb") as f: | |
| embed_metadata = pickle.load(f) | |
| embedder_model_name = embed_metadata["model_name"] | |
| embedder = SentenceTransformer(embedder_model_name, device=DEVICE) | |
| n_classes = len(label_maps["y6"]["classes"]) | |
| input_dim = int(embed_metadata["embedding_dim"]) | |
| model = FlatEmbedMLP( | |
| input_dim=input_dim, | |
| n_classes=n_classes, | |
| ).to(DEVICE) | |
| model_path = artifacts_dir / "models" / "flat_embed_best.pt" | |
| model.load_state_dict(torch.load(model_path, map_location=DEVICE)) | |
| model.eval() | |
| title_lookup_path = artifacts_dir / "label_maps" / "y6_title_lookup.json" | |
| with open(title_lookup_path, "r") as f: | |
| y6_title_lookup = json.load(f) | |
| _artifacts = { | |
| "device": DEVICE, | |
| "embedder": embedder, | |
| "model": model, | |
| "label_maps": label_maps, | |
| "embed_metadata": embed_metadata, | |
| "y6_title_lookup": y6_title_lookup, | |
| } | |
| return _artifacts |