App

app.py

@@ -0,0 +1,50 @@
+import torch
+import torchaudio
+from transformers import AutoFeatureExtractor, AutoModelForAudioClassification
+import gradio as gr
+import json
+
+
+config = json.load(open("config.json"))
+
+feature_extractor = AutoFeatureExtractor.from_pretrained(
+    config['extractor_path'], do_normalize=True, return_attention_mask=True
+)
+
+model = AutoModelForAudioClassification.from_pretrained(
+    config['model_path']
+)
+
+def audio_pipeline(audio_file):
+    if isinstance(audio_file, str):
+        waveform, sample_rate = torchaudio.load(audio_file)
+    else:
+        waveform, sample_rate = torchaudio.load(audio_file.name)
+    waveform = waveform.mean(dim=0)
+    if sample_rate != feature_extractor.sampling_rate:
+        transform = torchaudio.transforms.Resample(
+            orig_freq=sample_rate,
+            new_freq=feature_extractor.sampling_rate)
+        waveform = transform(waveform)
+    inputs = feature_extractor(waveform,
+        sampling_rate=feature_extractor.sampling_rate,
+        return_tensors="pt",
+        padding=True)
+    with torch.no_grad():
+        logits = model(**inputs).logits
+    probs = torch.nn.functional.softmax(logits, dim=-1)[0]
+    
+    top_probs, top_ids = torch.topk(probs, config['top_k'])
+    top_labels = [model.config.id2label[idx.item()] for idx in top_ids]
+    
+    results = {label: prob.item() for label, prob in zip(top_labels, top_probs)}
+    return results
+
+demo = gr.Interface(
+    fn=audio_pipeline,
+    inputs=[gr.Audio(type="filepath", label="Upload Audio")],
+    outputs=gr.Label(num_top_classes=config['top_k']),
+    title="Music Mind",
+)
+
+demo.launch(debug=True)

encode_dataset.py

@@ -1,34 +0,0 @@
-import numpy as np
-import librosa
-from transformers import AutoFeatureExtractor
-import argparse
-import os
-
-
-def parse_args():
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--input_dir", type=str, default="data")
-    ap.add_argument("--output_dir", type=str, default="features")
-    ap.add_argument("--model_id", type=str)
-
-    return vars(ap.parse_args())
-
-def main(args):
-    feature_extractor = AutoFeatureExtractor.from_pretrained(args["model_id"])
-    fe_sr = feature_extractor.sampling_rate
-
-    for root, _, files in os.walk(args["input_dir"]):
-        for file in files:
-            if not file.endswith(".wav"):
-                continue
-            path = os.path.join(root, file)
-            waveform, _ = librosa.load(path, sr=fe_sr)
-            if waveform.ndim == 1:
-                waveform = waveform.reshape(1, -1)
-            inputs = feature_extractor(waveform, sampling_rate=fe_sr)
-            output_path = os.path.join(args["output_dir"], file.replace(".wav", ".npy"))
-            os.makedirs(os.path.dirname(output_path), exist_ok=True)
-            np.save(output_path, inputs['input_values'])
-
-if __name__ == "__main__":
-    main(parse_args())

push_model.py

@@ -1,29 +0,0 @@
-import argparse
-from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig
-from huggingface_hub import HfApi, HfFolder, Repository
-
-def push_model_to_hub(local_model_path, model_name):
-    model = AutoModelForSequenceClassification.from_pretrained(local_model_path)
-    tokenizer = AutoTokenizer.from_pretrained(local_model_path)
-    config = AutoConfig.from_pretrained(local_model_path)
-
-    temp_dir = f"./{model_name}"
-    model.save_pretrained(temp_dir)
-    tokenizer.save_pretrained(temp_dir)
-    config.save_pretrained(temp_dir)
-
-    api = HfApi()
-    user = HfFolder.get_user()
-    repo_url = api.create_repo(name=model_name, token=HfFolder.get_token())
-    
-    repo = Repository(local_dir=temp_dir, clone_from=repo_url)
-    repo.push_to_hub(commit_message="Initial commit")
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Push a model to the Hugging Face Hub")
-    parser.add_argument("--local_model_path", type=str, required=True, help="Path to the local model")
-    parser.add_argument("--model_name", type=str, required=True, help="Desired name of the model on the Hugging Face Hub")
-
-    args = parser.parse_args()
-    
-    push_model_to_hub(args.local_model_path, args.model_name)

train.py

@@ -1,134 +0,0 @@
-from transformers import AutoModelForAudioClassification
-from torch.utils.data import DataLoader
-import evaluate
-import torch
-from tqdm import tqdm
-import argparse
-import json
-import os
-import shutil
-import mlflow
-import mlflow.pytorch
-
-from gtzan import GtzanDataset
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-metric = evaluate.load("accuracy")
-
-def parse_args():
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--label2id", type=str)
-    ap.add_argument("--model_id", type=str)
-    ap.add_argument("--batch_size", type=int, default=32)
-    ap.add_argument("--train_dir", type=str, default="data/train")
-    ap.add_argument("--val_dir", type=str, default="data/val")
-    ap.add_argument("--num_workers", type=int, default=4)
-    ap.add_argument("--lr", type=float, default=1e-4)
-    ap.add_argument("--epochs", type=int, default=10)
-    ap.add_argument("--output_dir", type=str, default="./weights")
-    ap.add_argument("--seed", type=int, default=42)
-    ap.add_argument("--name", type=str, default="model")
-    return vars(ap.parse_args())
-
-def train(args):
-    torch.manual_seed(args["seed"])
-    
-    label2id = json.load(open(args["label2id"]))
-    id2label = {v: k for k, v in label2id.items()}
-    num_labels = len(label2id)
-    if not os.path.exists(args["output_dir"]):
-        os.makedirs(args["output_dir"])
-
-    train_dataset = GtzanDataset(args["train_dir"], label2id)
-    val_dataset = GtzanDataset(args["val_dir"], label2id)
-
-    train_loader = DataLoader(
-        train_dataset,
-        batch_size=args["batch_size"],
-        shuffle=True,
-        num_workers=args["num_workers"])
-
-    val_loader = DataLoader(
-        val_dataset,
-        batch_size=args["batch_size"],
-        shuffle=False,
-        num_workers=args["num_workers"])
-
-    model = AutoModelForAudioClassification.from_pretrained(
-        args['model_id'],
-        num_labels=num_labels,
-        label2id=label2id,
-        id2label=id2label,
-    ).to(device)
-
-    optimizer = torch.optim.AdamW(model.parameters(), lr=args["lr"])
-    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
-        optimizer, T_max=len(train_loader) * args["epochs"]
-    )
-
-    max_val_accuracy = 0
-    best_path = ""
-
-    with mlflow.start_run():
-        mlflow.log_params({
-            "model_id": args["model_id"],
-            "batch_size": args["batch_size"],
-            "lr": args["lr"],
-            "epochs": args["epochs"],
-            "seed": args["seed"]
-        })
-
-        for epoch in tqdm(range(args["epochs"])):
-            model.train()
-            train_progress_bar = tqdm(train_loader, desc=f"Training Epoch {epoch + 1}")
-            for batch in train_progress_bar:
-                input_values, attention_mask, label = [b.to(device) for b in batch]
-                outputs = model(input_values=input_values,
-                                attention_mask=attention_mask,
-                                labels=label)
-                loss = outputs.loss
-                loss.backward()
-                optimizer.step()
-                lr_scheduler.step()
-                optimizer.zero_grad()
-
-                train_progress_bar.set_postfix({"loss": loss.item()})
-                train_progress_bar.update(1)
-                mlflow.log_metric("train_loss", loss.item())  # Log training loss
-            
-            torch.cuda.empty_cache()
-            model.eval()
-
-            val_progress_bar = tqdm(val_loader, desc="Validation")
-            for batch in val_progress_bar:
-                input_values, attention_mask, label = [b.to(device) for b in batch]
-                with torch.no_grad():
-                    outputs = model(input_values=input_values,
-                                    attention_mask=attention_mask,
-                                    labels=label)
-
-                logits = outputs.logits
-                predictions = torch.argmax(logits, dim=-1)
-                metric.add_batch(predictions=predictions, references=label)
-                val_progress_bar.update(1)
-            
-            val_accuracy = metric.compute()
-            mlflow.log_metric("val_accuracy", val_accuracy["accuracy"], step=epoch)  # Log validation accuracy
-            torch.cuda.empty_cache()
-            if val_accuracy["accuracy"] > max_val_accuracy:
-                if best_path:
-                    shutil.rmtree(best_path)
-                model_save_dir = os.path.join(
-                    args["output_dir"],
-                    args['name'],
-                    f"{int(round(val_accuracy['accuracy'], 2) * 100)}")
-                if not os.path.exists(model_save_dir):
-                    os.makedirs(model_save_dir, exist_ok=True)
-                model.save_pretrained(model_save_dir)
-                max_val_accuracy = val_accuracy["accuracy"]
-                best_path = model_save_dir
-        
-        mlflow.pytorch.log_model(model, "model")
-if __name__ == "__main__":
-    train(parse_args())