backend/src/dataset.py

"""Dataset loading and preprocessing for lyrical emotion analysis."""

import pandas as pd
from pathlib import Path
from typing import Optional
from datasets import Dataset, DatasetDict
from transformers import AutoTokenizer

from .config import config


class LyricsDataset:
    """Handler for lyrics emotion datasets."""

    def __init__(self, tokenizer_name: str = None):
        self.tokenizer = AutoTokenizer.from_pretrained(
            tokenizer_name or config.model_name
        )
        self.label2id = {label: i for i, label in enumerate(config.label_names)}
        self.id2label = {i: label for i, label in enumerate(config.label_names)}

    def load_from_csv(self, filepath: Path, lyrics_col: str = "lyrics",
                      label_col: str = "mood") -> DatasetDict:
        """Load dataset from CSV file."""
        df = pd.read_csv(filepath)

        # Normalize labels to lowercase
        df[label_col] = df[label_col].str.lower().str.strip()

        # Filter to valid labels only
        valid_mask = df[label_col].isin(self.label2id.keys())
        df = df[valid_mask].copy()

        # Map labels to IDs
        df["label"] = df[label_col].map(self.label2id)
        df["text"] = df[lyrics_col]

        # Create HuggingFace dataset
        dataset = Dataset.from_pandas(df[["text", "label"]])

        # Split into train/val/test (80/10/10)
        train_test = dataset.train_test_split(test_size=0.2, seed=42)
        val_test = train_test["test"].train_test_split(test_size=0.5, seed=42)

        return DatasetDict({
            "train": train_test["train"],
            "validation": val_test["train"],
            "test": val_test["test"]
        })

    def create_sample_dataset(self) -> DatasetDict:
        """Create a small sample dataset for testing the pipeline."""
        sample_data = {
            "text": [
                "I feel so alive today, the sun is shining and everything is perfect",
                "Dancing in the moonlight, feeling so free and happy",
                "Life is beautiful when you're with me, my heart is full of joy",
                "Celebrate good times, come on, let's have a party tonight",
                "My heart is broken, you left me alone in the dark",
                "Tears falling down, I can't stop crying over you",
                "Empty rooms and silent nights, missing you so much",
                "The pain won't go away, drowning in my sorrow",
                "I'm so mad at you, how could you do this to me",
                "Burning rage inside, I want to scream and shout",
                "You betrayed my trust, I'll never forgive you",
                "Fire in my veins, this anger consumes me whole",
                "Peaceful evening by the lake, watching the sunset fade",
                "Gentle breeze and quiet thoughts, finding my inner peace",
                "Floating on clouds, everything is calm and serene",
                "Soft rain falling, wrapped in warmth and comfort",
            ],
            "label": [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3]
        }

        dataset = Dataset.from_dict(sample_data)
        train_test = dataset.train_test_split(test_size=0.25, seed=42)

        return DatasetDict({
            "train": train_test["train"],
            "validation": train_test["test"],
            "test": train_test["test"]
        })

    def tokenize(self, dataset: DatasetDict) -> DatasetDict:
        """Tokenize the dataset."""
        def tokenize_fn(examples):
            return self.tokenizer(
                examples["text"],
                padding="max_length",
                truncation=True,
                max_length=config.max_length
            )

        tokenized = dataset.map(tokenize_fn, batched=True)
        tokenized = tokenized.remove_columns(["text"])
        tokenized.set_format("torch")

        return tokenized


def download_moodylyrics():
    """Instructions for obtaining MoodyLyrics dataset."""
    print("""
    ╔══════════════════════════════════════════════════════════════════╗
    ║                    MoodyLyrics Dataset Setup                     ║
    ╠══════════════════════════════════════════════════════════════════╣
    ║                                                                  ║
    ║  MoodyLyrics is a research dataset. To obtain it:                ║
    ║                                                                  ║
    ║  Option 1: Academic access                                       ║
    ║  - Request from the original authors                             ║
    ║  - Paper: "MoodyLyrics: A Sentiment Annotated Lyrics Dataset"    ║
    ║                                                                  ║
    ║  Option 2: Use similar Kaggle datasets                           ║
    ║  - Search for "lyrics emotion dataset" on Kaggle                 ║
    ║  - Download and save as data/lyrics_emotion.csv                  ║
    ║                                                                  ║
    ║  Expected CSV format:                                            ║
    ║  - Column 'lyrics': song lyrics text                             ║
    ║  - Column 'mood': one of [happy, sad, angry, relaxed]            ║
    ║                                                                  ║
    ║  For quick testing, use create_sample_dataset() method           ║
    ║                                                                  ║
    ╚══════════════════════════════════════════════════════════════════╝
    """)