File size: 6,196 Bytes
29f7657 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | #!/usr/bin/env python3
"""
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/*")
# private repo: run `huggingface-cli login` once, or set the HF_TOKEN environment variable
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 # model definitions (same code that trained the models)
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": # average the single-modality models
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()
|