Spaces:
Running on Zero
Running on Zero
| """Model wrappers (LightGBM, MDN) + the Engine that routes requests to them.""" | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import joblib | |
| import numpy as np | |
| import torch | |
| from .core import predict_velocities | |
| from ..models import heads | |
| from ..models.model import VelocityTransformer | |
| from ..data.seqdata import build_split_tensors, scatter_predictions | |
| class LgbmModel: | |
| def __init__(self, bundle): | |
| self.model = bundle["model"] | |
| self.cat = list(bundle["cat"]) | |
| self.drop = list(bundle["drop"]) | |
| self.cat_categories = bundle["cat_categories"] | |
| self.best_iteration = int(bundle["best_iteration"]) | |
| def predict_all(self, df): | |
| X = df.drop(columns=self.drop) | |
| for c in self.cat: | |
| X[c] = X[c].astype("category").cat.set_categories(self.cat_categories[c]) | |
| return np.asarray(self.model.predict(X, num_iteration=self.best_iteration), dtype=float) | |
| class TransformerModel: | |
| """Serves the shared transformer backbone with a sampling head. | |
| head_type selects the output head + sampler: "mdn" (temperature-controlled | |
| mixture) or "categorical" (softmax over velocity bins; temperature ignored | |
| by heads.sample). Both condition on genre only (not style). | |
| """ | |
| def __init__(self, transformer, genre_vocab, bpm_mean, bpm_std, device="cpu", head_type="mdn"): | |
| self.m = transformer.to(device).eval() | |
| self.genre_vocab = genre_vocab | |
| self.bpm_mean = float(bpm_mean) | |
| self.bpm_std = float(bpm_std) | |
| self.device = device | |
| self.head_type = head_type | |
| def load(cls, meta_path, ckpt_path, device="cpu", head_type="mdn"): | |
| with open(meta_path) as fh: | |
| meta = json.load(fh) | |
| gv = meta["genre_vocab"] | |
| m = VelocityTransformer(n_genres=len(gv) + 1, head=head_type) | |
| m.load_state_dict(torch.load(ckpt_path, map_location=device)["best_model"]) | |
| return cls(m, gv, meta["bpm_mean"], meta["bpm_std"], device, head_type) | |
| def predict_all(self, df, temperature=1.0, seed=42): | |
| t = build_split_tensors(df, self.genre_vocab, self.bpm_mean, self.bpm_std) | |
| gen = torch.Generator().manual_seed(int(seed)) | |
| with torch.no_grad(): | |
| raw = self.m(t["voice_idx"].to(self.device), t["genre_idx"].to(self.device), | |
| t["num_feats"].to(self.device), t["pad_mask"].to(self.device)).cpu() | |
| s = heads.sample(self.head_type, raw, generator=gen, temperature=float(temperature)) | |
| return scatter_predictions(t["row_idx"], s, t["pad_mask"], len(df)) | |
| MdnModel = TransformerModel # backward-compatible alias (head_type defaults to "mdn") | |
| class Engine: | |
| def __init__(self, lgbm, mdn, cat, styles, genres): | |
| self.lgbm = lgbm | |
| self.mdn = mdn | |
| self.cat = cat | |
| self.styles = list(styles) | |
| self.genres = list(genres) | |
| def load(cls, proc_dir=os.path.join("data", "processed")): | |
| lgbm = LgbmModel(joblib.load(os.path.join(proc_dir, "lightgbm_model.joblib"))) | |
| mdn = TransformerModel.load(os.path.join(proc_dir, "mdn_meta.json"), | |
| os.path.join(proc_dir, "head_mdn.pt"), head_type="mdn") | |
| cat = TransformerModel.load(os.path.join(proc_dir, "transformer_meta.json"), | |
| os.path.join(proc_dir, "head_categorical.pt"), | |
| head_type="categorical") | |
| with open(os.path.join(proc_dir, "lightgbm_features.json")) as fh: | |
| feats = json.load(fh) | |
| lv = feats["categorical_levels"] | |
| return cls(lgbm, mdn, cat, lv["style"], lv["genre"]) | |
| def predict(self, request): | |
| model = request["model"] | |
| seed = int(request.get("seed", 42)) | |
| if model == "lgbm": | |
| predict_all = self.lgbm.predict_all | |
| elif model == "mdn": | |
| temp = float(request.get("temperature", 1.0)) | |
| predict_all = lambda df: self.mdn.predict_all(df, temperature=temp, seed=seed) # noqa: E731 | |
| elif model == "categorical": | |
| predict_all = lambda df: self.cat.predict_all(df, seed=seed) # noqa: E731 (temperature n/a) | |
| else: | |
| raise ValueError(f"unknown model {model!r}") | |
| return predict_velocities(request, predict_all) | |
| def levels(self): | |
| return {"models": ["lgbm", "mdn", "categorical"], "styles": self.styles, "genres": self.genres} | |