src/data_loader.py
import os
import numpy as np
import pandas as pd
from torch.utils.data import Dataset, DataLoader
import librosa

class RedMusicDataset(Dataset):
    """
    Dataset expects:
    - lyrics csv: columns = ['resource_id', 'lyrics_text', ...]
    - audio files: <resource_id>.wav or .mp3 under audio_dir
    - interactions csv: ['user_id','resource_id','rating','timestamp', ...]
    """
    def __init__(self, lyrics_csv, interactions_csv, audio_dir, config, mode="train"):
        self.lyrics = pd.read_csv(lyrics_csv)
        self.interactions = pd.read_csv(interactions_csv)
        self.audio_dir = audio_dir
        self.sample_rate = config.get("dataset", {}).get("sample_rate", 44100)
        # merge on resource_id; interactions drive rows
        self.data = pd.merge(self.interactions, self.lyrics, on="resource_id", how="left")
        self.data = self.data.dropna(subset=['lyrics_text'])
        self.mode = mode

    def __len__(self):
        return len(self.data)

    def _load_audio(self, resource_id):
        for ext in ['.wav', '.mp3', '.flac']:
            path = os.path.join(self.audio_dir, f"{resource_id}{ext}")
            if os.path.exists(path):
                y, sr = librosa.load(path, sr=self.sample_rate)
                return y, sr
        # fallback: silence (1 second)
        return np.zeros(self.sample_rate, dtype=np.float32), self.sample_rate

    def __getitem__(self, idx):
        row = self.data.iloc[idx]
        resource_id = str(row['resource_id'])
        user_id = str(row['user_id'])
        rating = row.get('rating', 1.0)
        audio, sr = self._load_audio(resource_id)
        lyrics = str(row['lyrics_text'])
        sample = {
            "user_id": user_id,
            "resource_id": resource_id,
            "audio": audio.astype(np.float32),
            "sr": sr,
            "lyrics": lyrics,
            "rating": float(rating)
        }
        return sample

def collate_fn(batch):
    # returns raw batch; model/trainer handles padding and conversion
    return batch

def get_dataloader(lyrics_csv, interactions_csv, audio_dir, config, batch_size=32, shuffle=True, num_workers=4):
    ds = RedMusicDataset(lyrics_csv, interactions_csv, audio_dir, config)
    return DataLoader(ds, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, collate_fn=collate_fn)


src/feature_extraction.py
import numpy as np
import librosa
import torch
from transformers import BertTokenizer, BertModel

class AudioFeatureExtractor:
    def __init__(self, sample_rate=44100, n_mels=128, hop_length=512, n_fft=2048):
        self.sr = sample_rate
        self.n_mels = n_mels
        self.hop = hop_length
        self.n_fft = n_fft

    def stft_magnitude(self, y):
        S = np.abs(librosa.stft(y, n_fft=self.n_fft, hop_length=self.hop))
        return S

    def mel_spectrogram(self, y):
        mel = librosa.feature.melspectrogram(y, sr=self.sr, n_fft=self.n_fft, hop_length=self.hop, n_mels=self.n_mels)
        mel_db = librosa.power_to_db(mel, ref=np.max)
        return mel_db.astype(np.float32)

    def compute_mel_and_stft(self, y):
        mel = self.mel_spectrogram(y)
        stft = self.stft_magnitude(y)
        return mel, stft

    def time_frequency_attention(self, mel):
        # lightweight TF attention placeholder: elementwise weighting
        freq_attn = np.mean(mel, axis=1, keepdims=True)  # (n_mels,1)
        time_attn = np.mean(mel, axis=0, keepdims=True)  # (1,T)
        freq_w = freq_attn / (np.max(freq_attn) + 1e-9)
        time_w = time_attn / (np.max(time_attn) + 1e-9)
        fused = mel * (freq_w @ time_w)
        return fused

class TextFeatureExtractor:
    def __init__(self, model_name="bert-base-chinese", device="cpu"):
        self.tokenizer = BertTokenizer.from_pretrained(model_name)
        self.model = BertModel.from_pretrained(model_name).to(device)
        self.device = device

    def encode(self, texts, max_length=128):
        enc = self.tokenizer(texts, padding=True, truncation=True, max_length=max_length, return_tensors="pt")
        input_ids = enc["input_ids"].to(self.device)
        attention_mask = enc["attention_mask"].to(self.device)
        with torch.no_grad():
            outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
            cls = outputs.last_hidden_state[:,0,:]
        return cls

src/models.py
import torch
import torch.nn as nn
import torch.nn.functional as F

class AudioCNN(nn.Module):
    """
    Small CNN backbone for mel-spectrogram input.
    Input shape: (batch, 1, n_mels, T)
    """
    def __init__(self, n_mels=128, embed_dim=128):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=(3,3), padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d((2,2)),
            nn.Conv2d(32, 64, kernel_size=(3,3), padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1,1))
        )
        self.fc = nn.Linear(64, embed_dim)

    def forward(self, x):
        h = self.conv(x)
        h = h.view(h.size(0), -1)
        return self.fc(h)

class CrossModalGating(nn.Module):
    """
    Simple gating fusion between audio and text embeddings.
    """
    def __init__(self, dim):
        super().__init__()
        self.fc_audio = nn.Linear(dim, dim)
        self.fc_text = nn.Linear(dim, dim)
        self.sigmoid = nn.Sigmoid()

    def forward(self, audio_emb, text_emb):
        a = torch.tanh(self.fc_audio(audio_emb))
        t = torch.tanh(self.fc_text(text_emb))
        gate = self.sigmoid(a + t)
        fused = gate * audio_emb + (1 - gate) * text_emb
        return fused

class RecommenderHead(nn.Module):
    """
    Small recommendation head. User embedding uses a placeholder large embedding table.
    Replace with a proper mapping/dictionary for large-scale datasets.
    """
    def __init__(self, embed_dim=128, num_users=1000000):
        super().__init__()
        self.user_embeddings = nn.Embedding(num_users, embed_dim)
        self.resource_fc = nn.Linear(embed_dim, embed_dim)
        self.predictor = nn.Sequential(
            nn.Linear(embed_dim * 2, embed_dim),
            nn.ReLU(),
            nn.Linear(embed_dim, 1)
        )

    def forward(self, user_ids, resource_emb):
        # user_ids: list or tensor
        if isinstance(user_ids, list):
            user_ids = torch.tensor([int(u) for u in user_ids], device=resource_emb.device, dtype=torch.long)
        user_vec = self.user_embeddings(user_ids)
        res_vec = self.resource_fc(resource_emb)
        x = torch.cat([user_vec, res_vec], dim=-1)
        score = self.predictor(x).squeeze(-1)
        return score

class RedMusicRecommender(nn.Module):
    def __init__(self, config):
        super().__init__()
        embed_dim = config.get("model", {}).get("embed_dim", 128)
        n_mels = config.get("model", {}).get("audio_n_mels", 128)
        self.audio_net = AudioCNN(n_mels=n_mels, embed_dim=embed_dim)
        self.fusion = CrossModalGating(dim=embed_dim)
        self.recommender = RecommenderHead(embed_dim=embed_dim)

    def forward(self, audio_batch, text_embeddings, user_ids):
        audio_emb = self.audio_net(audio_batch)
        fused = self.fusion(audio_emb, text_embeddings)
        scores = self.recommender(user_ids, fused)
        return scores, fused


src/trainer.py
import os
import numpy as np
import torch
import pytorch_lightning as pl
from torch.utils.data import DataLoader
import mlflow
from .models import RedMusicRecommender
from .feature_extraction import AudioFeatureExtractor, TextFeatureExtractor
from .data_loader import get_dataloader
from .utils import load_config

class PLModule(pl.LightningModule):
    def __init__(self, config):
        super().__init__()
        self.save_hyperparameters()
        self.config = config
        self.model = RedMusicRecommender(config)
        self.loss_fn = torch.nn.BCEWithLogitsLoss()
        self.lr = config.get("training", {}).get("lr", 1e-3)

    def forward(self, *args, **kwargs):
        return self.model(*args, **kwargs)

    def training_step(self, batch, batch_idx):
        # batch is a list of samples from collate_fn
        audio_list = []
        text_list = []
        user_ids = []
        labels = []
        afe = AudioFeatureExtractor(sample_rate=self.config['dataset']['sample_rate'],
                                    n_mels=self.config['model']['audio_n_mels'])
        for s in batch:
            mel, _ = afe.compute_mel_and_stft(s['audio'])
            T_max = 128
            if mel.shape[1] < T_max:
                pad = np.zeros((mel.shape[0], T_max - mel.shape[1]), dtype=np.float32)
                mel = np.concatenate([mel, pad], axis=1)
            else:
                mel = mel[:, :T_max]
            audio_list.append(mel)
            text_list.append(s['lyrics'])
            user_ids.append(s['user_id'])
            labels.append(1.0 if s.get('rating', 1.0) >= 1.0 else 0.0)

        audios = torch.from_numpy(np.stack(audio_list)).unsqueeze(1).to(self.device)  # (B,1,n_mels,T)
        tfe = TextFeatureExtractor(device=self.device)
        text_embs = tfe.encode(text_list)  # (B, hidden)
        labels = torch.tensor(labels, dtype=torch.float32, device=self.device)

        scores, _ = self.model(audios, text_embs, user_ids)
        loss = self.loss_fn(scores, labels)
        self.log("train_loss", loss)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr, weight_decay=1e-4)
        return optimizer

def train_main(config_path="config.yaml"):
    config = load_config(config_path)
    dataloader = get_dataloader(config['dataset']['lyrics_file'],
                                config['dataset']['interactions_file'],
                                config['dataset']['audio_dir'],
                                config,
                                batch_size=config['training']['batch_size'])
    pl.seed_everything(42)
    model = PLModule(config)
    trainer = pl.Trainer(accelerator="auto", devices=1 if torch.cuda.is_available() else None,
                         max_epochs=config['training']['epochs'])
    mlflow.start_run()
    trainer.fit(model, dataloader)
    mlflow.end_run()
    return model

if __name__ == "__main__":
    train_main()

src/inference.py
import torch
import numpy as np
from .models import RedMusicRecommender
from .feature_extraction import AudioFeatureExtractor, TextFeatureExtractor
from .utils import load_config

def load_checkpoint(checkpoint_path, config):
    model = RedMusicRecommender(config)
    state = torch.load(checkpoint_path, map_location="cpu")
    if 'state_dict' in state:
        model.load_state_dict(state['state_dict'])
    else:
        model.load_state_dict(state)
    model.eval()
    return model

def recommend(model, samples, top_k=10, device='cpu'):
    afe = AudioFeatureExtractor(sample_rate=44100, n_mels=128)
    tfe = TextFeatureExtractor(device=device)
    audios = []
    texts = []
    user_ids = []
    for s in samples:
        mel, _ = afe.compute_mel_and_stft(s['audio'])
        T_max = 128
        if mel.shape[1] < T_max:
            pad = np.zeros((mel.shape[0], T_max - mel.shape[1]), dtype=np.float32)
            mel = np.concatenate([mel, pad], axis=1)
        else:
            mel = mel[:, :T_max]
        audios.append(mel)
        texts.append(s['lyrics'])
        user_ids.append(s['user_id'])
    audios = torch.from_numpy(np.stack(audios)).unsqueeze(1).to(device)
    text_embs = tfe.encode(texts)
    scores, fused = model(audios, text_embs, user_ids)
    scores = scores.detach().cpu().numpy()
    top_idx = np.argsort(-scores)[:top_k]
    return [(int(i), float(scores[i])) for i in top_idx]

if __name__ == "__main__":
    # Example inference usage: load config, checkpoint, and run recommend() with sample batch
    cfg = load_config("config.yaml")
    # model = load_checkpoint("path/to/checkpoint.pth", cfg)
    # samples = [{ "audio": np.zeros(cfg['dataset']['sample_rate']), "lyrics": "sample lyrics", "user_id": "1" }]
    # print(recommend(model, samples))
    pass


src/evaluation.py
import numpy as np
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score

def compute_metrics(y_true, y_scores, threshold=0.5):
    y_pred = (np.array(y_scores) >= threshold).astype(int)
    y_true = np.array(y_true).astype(int)
    return {
        "accuracy": float(accuracy_score(y_true, y_pred)),
        "precision": float(precision_score(y_true, y_pred, zero_division=0)),
        "recall": float(recall_score(y_true, y_pred, zero_division=0)),
        "f1": float(f1_score(y_true, y_pred, zero_division=0))
    }


src/utils.py
import yaml

def load_config(path="config.yaml"):
    with open(path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)