| """LAION-CLAP's music checkpoint on the same benchmark as song_describer.py, for comparison. |
| |
| uv run clap_eval.py |
| """ |
| import argparse |
| from pathlib import Path |
|
|
| import laion_clap |
| import numpy as np |
| import torch |
| from huggingface_hub import hf_hub_download |
|
|
| from song_describer import clips, load_captions, load_tracks, print_results |
|
|
| SR = 48000 |
| CHECKPOINT = "music_audioset_epoch_15_esc_90.14.pt" |
|
|
|
|
| def normalize(vectors): |
| return vectors / np.linalg.norm(vectors, axis=-1, keepdims=True) |
|
|
|
|
| @torch.inference_mode() |
| def embed(model, wave): |
| size = 10 * SR |
| chunks = np.stack([wave[start:start + size] for start in range(0, len(wave) - size + 1, size)]) |
| |
| chunks = (np.clip(chunks, -1, 1) * 32767).astype(np.int16).astype(np.float32) / 32767 |
| return normalize(model.get_audio_embedding_from_data(chunks).mean(0)) |
|
|
|
|
| @torch.inference_mode() |
| def embed_text(model, captions): |
| return normalize(np.concatenate([model.get_text_embedding(captions[start:start + 64]) for start in range(0, len(captions), 64)])) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument("--data", type=Path, default=Path("data/song_describer")) |
| parser.add_argument("--device", default="cuda", help="cuda, mps or cpu") |
| args = parser.parse_args() |
|
|
| captions = load_captions(args.data) |
| model = laion_clap.CLAP_Module(enable_fusion=False, amodel="HTSAT-base", device=args.device) |
| model.load_ckpt(hf_hub_download("lukewys/laion_clap", CHECKPOINT), verbose=False) |
|
|
| tracks, middle, average = [], [], [] |
| for track, wave in load_tracks(args.data, SR): |
| three = [embed(model, clip) for clip in clips(wave, SR)] |
| tracks.append(track) |
| middle.append(three[1]) |
| average.append(normalize(np.mean(three, axis=0))) |
|
|
| order = np.argsort(tracks) |
| views = {"one 10 s clip": middle, "three 10 s clips": average} |
| views = {name: np.stack(audio)[order] for name, audio in views.items()} |
| text = embed_text(model, captions.text.tolist()) |
| print_results("CLAP music", text, views, captions["corpus-id"].to_numpy(), np.array(tracks)[order]) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
