| |
| """ |
| Predict the theme of Hindi songs (romance / sadness / celebration / motivational) with the models |
| of meet5568/meet_models. |
| |
| For a NEW song, every one of the 5 cross-validation folds has its own lyrics encoder and its own |
| classifier; the script runs all 5 and averages their probabilities. |
| |
| Examples |
| # lyrics only (downloads the models from Hugging Face the first time; needs access to the repo) |
| python predict.py --lyrics "tujhe dekha toh yeh jaana sanam, pyaar hota hai deewana sanam" |
| python predict.py --lyrics-file song.txt |
| # use an already downloaded copy of the repo |
| python predict.py --model-dir /path/to/meet_models --lyrics-file song.txt |
| # lyrics + audio: an .npz with arrays "instrument" and "mix" (and/or "vocal"), each (n_songs, 9984) |
| python predict.py --lyrics-file song.txt --audio-npz song_audio.npz --combination lyrics+instrument+mix |
| |
| Python API |
| from predict import ThemePredictor |
| p = ThemePredictor() # or ThemePredictor(model_dir="...") |
| p.predict(["lyrics of song 1", "lyrics of song 2"]) # DataFrame of probabilities |
| p.predict(lyrics, audio={"instrument": A, "mix": B}, combination="lyrics+instrument+mix") |
| """ |
|
|
| import argparse |
| import json |
| import os |
| import sys |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
|
|
| REPO = "meet5568/meet_models" |
| THEMES = ["romance", "sadness", "celebration", "motivational"] |
| AUDIO = ("vocal", "instrument", "mix") |
|
|
|
|
| def get_model_dir(model_dir=None, need_encoders=True, need_classifiers=True): |
| """A local copy of the repo; downloads (only the needed parts) if no path is given.""" |
| if model_dir: |
| return model_dir |
| from huggingface_hub import snapshot_download |
| patterns = ["code/*", "README.md"] |
| if need_encoders: |
| patterns.append("lyrics_encoder/theme_tuned/*") |
| if need_classifiers: |
| patterns.append("classification/*") |
| |
| return snapshot_download(REPO, allow_patterns=patterns, token=os.environ.get("HF_TOKEN")) |
|
|
|
|
| class ThemePredictor: |
| def __init__(self, model_dir=None, device=None): |
| self.dir = get_model_dir(model_dir) |
| sys.path.insert(0, os.path.join(self.dir, "code")) |
| import classify |
| self.C = classify |
| self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| def encode_lyrics(self, texts, fold): |
| from sentence_transformers import SentenceTransformer |
| enc = SentenceTransformer(os.path.join(self.dir, "lyrics_encoder", "theme_tuned", f"fold_{fold}"), |
| device=self.device) |
| enc.max_seq_length = 400 |
| return enc.encode(texts, batch_size=16, normalize_embeddings=True, convert_to_numpy=True) |
|
|
| def _fold_proba(self, combination, fold, inputs): |
| d = os.path.join(self.dir, "classification", combination, f"fold_{fold}") |
| info = json.load(open(os.path.join(d, "config.json"))) |
| if info["method"] == "late": |
| return np.mean([self._fold_proba(m, fold, inputs) for m in info["modalities"]], axis=0) |
| norm = np.load(os.path.join(d, "norm.npz")) |
| model = self.C.build_model(info["run"], info["modalities"], info["in_dims"], info["model_config"]) |
| model.load_state_dict(torch.load(os.path.join(d, "model.pt"), map_location="cpu")) |
| model = model.to(self.device).eval() |
| xs = [torch.tensor((inputs[m] - norm[f"{m}_mean"]) / norm[f"{m}_std"], dtype=torch.float32, |
| device=self.device) for m in info["modalities"]] |
| with torch.no_grad(): |
| return self.C.predict_proba(model, xs) |
|
|
| def predict(self, lyrics=None, audio=None, combination=None): |
| """lyrics: list of strings (or None); audio: dict modality -> (n, 9984) array (or None). |
| combination: e.g. "lyrics", "lyrics+instrument+mix", "mix"; default = all given inputs.""" |
| audio = {k: np.asarray(v, dtype=np.float32) for k, v in (audio or {}).items()} |
| if lyrics is not None and isinstance(lyrics, str): |
| lyrics = [lyrics] |
| given = (["lyrics"] if lyrics is not None else []) + [m for m in AUDIO if m in audio] |
| combination = combination or "+".join(given) |
| mods = combination.split("+") |
| missing = [m for m in mods if m not in given] |
| if missing: |
| raise ValueError(f"combination {combination} needs inputs {missing}") |
| n = len(lyrics) if lyrics is not None else len(next(iter(audio.values()))) |
| probs = [] |
| for k in range(5): |
| inputs = dict(audio) |
| if "lyrics" in mods: |
| inputs["lyrics"] = self.encode_lyrics(lyrics, k) |
| probs.append(self._fold_proba(combination, k, inputs)) |
| p = np.mean(probs, axis=0) |
| out = pd.DataFrame(p, columns=THEMES) |
| out.insert(0, "predicted_theme", [THEMES[i] for i in p.argmax(1)]) |
| assert len(out) == n |
| return out |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description="Predict the theme of Hindi songs") |
| ap.add_argument("--lyrics", help="lyrics text of one song") |
| ap.add_argument("--lyrics-file", help="text file with the lyrics of one song") |
| ap.add_argument("--audio-npz", help=".npz with arrays vocal / instrument / mix, each (n_songs, 9984)") |
| ap.add_argument("--combination", default=None, help="e.g. lyrics, lyrics+instrument+mix (default: all inputs given)") |
| ap.add_argument("--model-dir", default=None, help="local copy of meet5568/meet_models (default: download)") |
| args = ap.parse_args() |
| lyrics = args.lyrics or (open(args.lyrics_file, encoding="utf-8").read() if args.lyrics_file else None) |
| audio = dict(np.load(args.audio_npz)) if args.audio_npz else None |
| if lyrics is None and audio is None: |
| ap.error("give --lyrics / --lyrics-file and/or --audio-npz") |
| p = ThemePredictor(args.model_dir).predict([lyrics] if lyrics else None, audio, args.combination) |
| print(p.round(3).to_string(index=False)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
