Audio-Tune: Fine-Tuning AST and RAG Creation

This project fine-tunes the MIT Audio Spectrogram Transformer (AST) model on a custom music dataset and creates a Retrieval-Augmented Generation (RAG) system using ChromaDB for semantic music search.

Overview

The workflow consists of three main stages:

1. Data Preparation: Scan your music library and create training/validation splits
2. Model Fine-Tuning: Fine-tune the AST model on your music styles
3. RAG Creation: Generate embeddings using the fine-tuned model and store them in ChromaDB

Prerequisites

Hardware Requirements

• CUDA-compatible GPU (or Intel XPU for Arc GPUs)
• At least 16GB RAM
• Sufficient disk space for your music library and model checkpoints

Software Requirements

• Python 3.8+
• CUDA toolkit (if using NVIDIA GPU) or Intel Extension for PyTorch (if using Intel Arc GPU)

Install Dependencies

pip install -r requirements.txt

The main dependencies include:

• torch - PyTorch framework
• transformers - Hugging Face transformers
• librosa - Audio processing
• chromadb - Vector database for embeddings
• plexapi - Plex server integration
• datasets - Hugging Face datasets
• evaluate - Model evaluation metrics
• python-dotenv - Environment variable management

Configuration

Plex Server Setup

Create a .env file in the project root with your Plex credentials:

PLEX_URL=http://your-plex-server:32400
PLEX_TOKEN=your-plex-token

To get your Plex token, follow this guide.

Step 1: Data Preparation

1.1 Scan Your Music Library

First, scan your Plex music library to create a metadata CSV. This script fetches tracks from your Plex server and creates a dataset based on style tags:

python scanner.py --output_csv music_metadata.csv

What it does:

• Connects to your Plex server
• Retrieves all tracks with style tags
• Filters styles with at least 100 samples
• Creates music_metadata.csv with columns: file_path, genre

1.2 Create Training Splits

Split the metadata into training and validation sets:

python create_splits.py

What it does:

• Loads music_metadata.csv
• Creates integer labels for each genre
• Splits data 90/10 (train/validation) with stratification
• Saves train.csv and val.csv

Output files:

• train.csv - Training dataset with columns: file_path, genre, label
• val.csv - Validation dataset with same structure

Step 2: Model Fine-Tuning

Fine-tune the pre-trained AST model on your music styles:

python fine_tune_ast.py

Configuration (edit in script):

MODEL_ID = "MIT/ast-finetuned-audioset-10-10-0.4593"  # Base model
TARGET_SAMPLE_RATE = 16000                              # Audio sample rate
MAX_DURATION_S = 10.0                                   # Clip duration
OUTPUT_DIR = "./ast-finetuned-all-styles"              # Output directory
BATCH_SIZE = 14                                         # Adjust for GPU memory
NUM_EPOCHS = 3                                          # Training epochs
LEARNING_RATE = 3e-5                                    # Learning rate

What it does:

1. Loads training and validation datasets from CSV files
2. Preprocesses audio files:
   • Loads audio at 16kHz
   • Trims silence (60dB threshold)
   • Randomly samples 10-second clips from each track
   • Converts to mel spectrograms
3. Fine-tunes the AST model on your genres/styles
4. Saves the best model based on validation accuracy

Output:

• ./ast-finetuned-all-styles/ - Directory containing:
  • config.json - Model configuration
  • model.safetensors - Fine-tuned model weights
  • preprocessor_config.json - Audio preprocessor config
  • checkpoint-*/ - Training checkpoints

Training Tips:

• Reduce BATCH_SIZE if you encounter CUDA out-of-memory errors
• Training time depends on dataset size (expect several hours for large libraries)
• Monitor accuracy metrics to avoid overfitting
• Best model is automatically saved based on validation accuracy

Step 3: RAG Creation (Generate Embeddings)

Generate embeddings from your fine-tuned model and store them in ChromaDB:

python get_embeddings.py

Configuration (edit in script):

MODEL_PATH = "./ast-finetuned-all-styles"          # Path to fine-tuned model
TARGET_SAMPLE_RATE = 16000                          # Audio sample rate
CHUNK_DURATION_S = 10.0                             # Embedding chunk size
DB_PATH = "./chroma_db_silence_all_styles"         # ChromaDB path
COLLECTION_NAME = "audio_embeddings"                # Collection name

What it does:

1. Loads your fine-tuned model (base transformer, not classification head)
2. Connects to your Plex server
3. Fetches tracks added in the last 30 days (configurable)
4. For each track:
   • Loads and trims silence
   • Extracts 3 embeddings (start, middle, end)
   • Stores embeddings in ChromaDB with metadata:
     • Title, artist, album, year
     • Genres and styles
     • Plex rating key
     • File path
     • Segment position

Output:

• ./chroma_db_silence_all_styles/ - ChromaDB database directory
  • chroma.sqlite3 - SQLite database
  • Vector embeddings stored internally

Note: The script skips tracks already in the database, so you can run it incrementally as you add new music.

Using the Embeddings

Once embeddings are generated, you can use them for:

Semantic Music Search

Use embeddings_lookup.py to query similar tracks:

from embeddings_lookup import EmbeddingLookup

lookup = EmbeddingLookup()
similar_tracks = lookup.search_similar("path/to/song.mp3", n_results=10)

Music Visualization

Generate UMAP visualizations of your music library:

python myumap.py

This creates an interactive HTML visualization showing your music library in 2D space.

TUI (Text User Interface)

Launch the interactive terminal UI:

python tui.py

Project Structure

audio-tune/
├── README.md                          # This file
├── requirements.txt                   # Python dependencies
├── .env                               # Plex credentials (create this)
│
├── scanner.py                         # Scan Plex library
├── create_splits.py                   # Create train/val splits
├── fine_tune_ast.py                   # Fine-tune AST model
├── get_embeddings.py                  # Generate embeddings
│
├── embeddings_lookup.py               # Embedding query utilities
├── myumap.py                          # UMAP visualization
├── tui.py                             # Terminal UI
├── plex_actions.py                    # Plex API actions
├── cluster_subgenres.py               # Clustering utilities
│
├── music_metadata.csv                 # Generated metadata
├── train.csv                          # Training split
├── val.csv                            # Validation split
├── clusters.csv                       # Clustering results
│
├── ast-finetuned-all-styles/          # Fine-tuned model
└── chroma_db_silence_all_styles/      # ChromaDB database

Customization

Adjusting Fine-Tuning Parameters

Edit fine_tune_ast.py:

• Sample Rate: Change TARGET_SAMPLE_RATE (default: 16000)
• Clip Duration: Change MAX_DURATION_S (default: 10 seconds)
• Batch Size: Adjust BATCH_SIZE based on GPU memory
• Epochs: Increase NUM_EPOCHS for better convergence
• Learning Rate: Tune LEARNING_RATE (default: 3e-5)

Changing Embedding Strategy

Edit get_embeddings.py:

• Time Window: Modify track filtering (default: last 30 days)
• Segments: Change which segments to extract (start/middle/end)
• Chunk Duration: Adjust CHUNK_DURATION_S
• Silence Trimming: Modify top_db parameter in librosa.effects.trim()

Genre Filtering

Edit scanner.py:

• Minimum Samples: Change the threshold for genre inclusion (default: 100)
• Genre Consolidation: Add genre mapping rules
• Metadata Source: Use genres instead of styles

Troubleshooting

CUDA Out of Memory

• Reduce BATCH_SIZE in fine_tune_ast.py
• Reduce MAX_DURATION_S to process shorter clips
• Close other GPU-intensive applications

Audio Loading Errors

• Ensure librosa is properly installed
• Check that audio files are readable
• Verify file paths in CSV files are absolute paths

Plex Connection Issues

• Verify .env file exists with correct credentials
• Check Plex server is running and accessible
• Test connection with plexapi directly

Low Fine-Tuning Accuracy

• Increase NUM_EPOCHS
• Verify genre labels are consistent
• Check for class imbalance in your dataset
• Add more training data for underrepresented genres

Slow Embedding Generation

• Ensure model is loaded on GPU (.to("xpu") or .to("cuda"))
• Process in smaller batches
• Skip tracks that are already embedded

Performance Notes

• Fine-tuning time: ~2-6 hours for 1000-5000 songs (depends on hardware)
• Embedding generation: ~1-3 seconds per track (GPU-accelerated)
• Database queries: <100ms for similarity search with 10k+ embeddings

License

This project uses the MIT Audio Spectrogram Transformer model, which is licensed under BSD-3-Clause.

Acknowledgments

• MIT for the Audio Spectrogram Transformer (AST)
• Hugging Face for the transformers library
• ChromaDB for the vector database