Upload 3 files

attribution.py

@@ -0,0 +1,118 @@
+import os
+import torch
+import numpy as np
+import laion_clap
+import pandas as pd
+
+# Set device to GPU if available
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+def compute_clap_embeddings(audio_dir, clap_model, batch_size=32):
+    # Collect all audio file paths in the directory
+    audio_files = [os.path.join(audio_dir, fn) for fn in os.listdir(audio_dir)
+                   if os.path.isfile(os.path.join(audio_dir, fn))]
+    
+    # Add debug print
+    print(f"Found {len(audio_files)} files in {audio_dir}")
+    
+    if not audio_files:
+        print(f"No files found in directory: {audio_dir}")
+        return [], None
+
+    embeddings_list = []
+    filenames_list = []
+
+    # Process audio files in batches
+    for i in range(0, len(audio_files), batch_size):
+        batch_files = audio_files[i:i + batch_size]
+        with torch.no_grad():
+            try:
+                # Get embeddings for the batch
+                embeddings = clap_model.get_audio_embedding_from_filelist(x=batch_files, use_tensor=True)
+                embeddings_list.append(embeddings)
+                filenames_list.extend([os.path.basename(f) for f in batch_files])
+            except Exception as e:
+                print(f"Error processing batch starting at index {i}: {str(e)}")
+                print(f"Problematic files: {batch_files}")
+                continue
+
+    if not embeddings_list:
+        print("No embeddings were generated")
+        return [], None
+
+    # Concatenate all embeddings
+    all_embeddings = torch.cat(embeddings_list, dim=0)
+    return filenames_list, all_embeddings
+
+# Load CLAP model
+clap_model = laion_clap.CLAP_Module(enable_fusion=True, device=device)
+clap_model.load_ckpt()
+clap_model.eval()
+
+# Step 1: Compute embeddings for training tracks
+training_dir = "./training"
+training_filenames, training_embeddings = compute_clap_embeddings(training_dir, clap_model)
+training_embeddings = training_embeddings.to(device)
+
+# Step 2: Compute embeddings for test tracks
+test_dir = './test'
+test_filenames, test_embeddings = compute_clap_embeddings(test_dir, clap_model)
+
+# Normalize training embeddings
+training_embeddings = torch.nn.functional.normalize(training_embeddings, dim=1)
+
+if test_filenames:
+    test_embeddings = test_embeddings.to(device)
+    test_embeddings = torch.nn.functional.normalize(test_embeddings, dim=1)
+    
+    # Compute similarity matrix (test samples x training samples)
+    similarity_matrix = test_embeddings @ training_embeddings.T  # Shape: (num_test_samples, num_training_samples)
+    
+    # Convert similarities to attribution scores by normalizing
+    attribution_scores = similarity_matrix / similarity_matrix.sum(dim=1, keepdim=True)
+    
+    # Map filenames to attribution scores
+    attribution_dict = {}
+    for i, test_file in enumerate(test_filenames):
+        scores = attribution_scores[i].cpu().numpy()
+        attribution_dict[test_file] = dict(zip(training_filenames, scores))
+    
+    # Optional: Save attribution scores to a JSON file
+    import json
+    with open('attribution_scores.json', 'w') as f:
+        json.dump(attribution_dict, f, indent=4)
+else:
+    print("No test files found in the directory. Skipping test embeddings computation and similarity calculations.")
+
+# Function to save embeddings to CSV
+def save_embeddings_to_csv(filenames, embeddings, csv_filename):
+    # Convert embeddings to a list
+    embeddings_list = embeddings.cpu().numpy().tolist()
+    # Create a DataFrame
+    df = pd.DataFrame({
+        'filename': filenames,
+        'embedding': embeddings_list
+    })
+    # Save to CSV
+    df.to_csv(csv_filename, index=False)
+
+# Save training embeddings
+save_embeddings_to_csv(training_filenames, training_embeddings, 'training_embeddings.csv')
+
+# Optional: Save test embeddings if needed
+# save_embeddings_to_csv(test_filenames, test_embeddings, 'test_embeddings.csv')
+
+# Function to save embeddings and filenames
+def save_embeddings(filenames, embeddings, filename_prefix):
+    # Save embeddings
+    np.save(f'{filename_prefix}_embeddings.npy', embeddings.cpu().numpy())
+    # Save filenames
+    with open(f'{filename_prefix}_filenames.txt', 'w') as f:
+        for item in filenames:
+            f.write("%s\n" % item)
+
+# Save training embeddings
+save_embeddings(training_filenames, training_embeddings, 'training')
+
+# Optional: Save test embeddings if needed
+# save_embeddings(test_filenames, test_embeddings, 'test')

similarity.py

@@ -0,0 +1,122 @@
+import os
+import torch
+import laion_clap
+import pandas as pd
+import ast  # To convert string representation of list back to list
+import numpy as np
+
+class CLAPSimilarity:
+    def __init__(self, training_embeddings_prefix='training', clap_model=None, device=None):
+        if device is None:
+            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.device = device
+
+        if clap_model is None:
+            # Load CLAP model
+            self.clap_model = laion_clap.CLAP_Module(enable_fusion=True, device=self.device)
+            self.clap_model.load_ckpt()
+            self.clap_model.eval()
+        else:
+            self.clap_model = clap_model
+
+        # Load precomputed training embeddings from files
+        self.training_embeddings, self.training_filenames = self.load_embeddings(training_embeddings_prefix)
+
+        # Normalize training embeddings
+        self.training_embeddings = torch.nn.functional.normalize(self.training_embeddings, dim=1)
+
+    def load_embeddings(self, filename_prefix):
+        # Load embeddings
+        embeddings = np.load(f'{filename_prefix}_embeddings.npy')
+        embeddings = torch.tensor(embeddings, device=self.device)
+
+        # Load filenames
+        with open(f'{filename_prefix}_filenames.txt', 'r') as f:
+            filenames = [line.strip() for line in f]
+
+        return embeddings, filenames
+
+    def compute_similarity(self, input_data, input_type='audio', max_tracks=0):
+        """
+        Compute similarity scores between input data and training embeddings.
+
+        Parameters:
+        - input_data: Either a string (text prompt or path to audio file) or a list of strings.
+        - input_type: 'audio' or 'text'
+        - max_tracks: Maximum number of tracks to include in the results. 0 means all tracks.
+
+        Returns:
+        - similarity_scores: A dictionary mapping training filenames to normalized similarity scores.
+        """
+        with torch.no_grad():
+            if input_type == 'audio':
+                # If input_data is a path to an audio file
+                if isinstance(input_data, str):
+                    input_files = [input_data]
+                else:
+                    input_files = input_data
+                embeddings = self.clap_model.get_audio_embedding_from_filelist(
+                    x=input_files, use_tensor=True
+                ).to(self.device)
+            elif input_type == 'text':
+                # If input_data is a text string or list of strings
+                if isinstance(input_data, str):
+                    input_texts = [input_data]
+                else:
+                    input_texts = input_data
+                embeddings = self.clap_model.get_text_embedding(
+                    input_texts, use_tensor=True
+                ).to(self.device)
+            else:
+                raise ValueError("input_type must be 'audio' or 'text'")
+
+
+            # Normalize embeddings
+            embeddings = torch.nn.functional.normalize(embeddings, dim=1)
+
+            # Compute similarity scores
+            similarity_matrix = embeddings @ self.training_embeddings.T  # (input_samples, training_samples)
+
+            # For single input, process accordingly
+            if similarity_matrix.shape[0] == 1:
+                similarities = similarity_matrix[0]
+                similarities = similarities.cpu().numpy()
+                # Shift to positive values
+                similarities = similarities - similarities.min()
+                # Normalize scores to sum to 100
+                total = similarities.sum()
+                if total > 0:
+                    normalized_scores = (similarities / total) * 100
+                else:
+                    normalized_scores = similarities
+                # Create a dictionary of filenames and scores
+                similarity_scores = dict(zip(self.training_filenames, normalized_scores))
+                # Sort the scores in descending order
+                similarity_scores = dict(sorted(similarity_scores.items(), key=lambda item: item[1], reverse=True))
+                # Limit the number of tracks if max_tracks is specified
+                if max_tracks > 0:
+                    similarity_scores = dict(list(similarity_scores.items())[:max_tracks])
+            else:
+                # For multiple inputs, return a list of dictionaries
+                similarity_scores = []
+                for i in range(similarity_matrix.shape[0]):
+                    similarities = similarity_matrix[i]
+                    similarities = similarities.cpu().numpy()
+                    # Shift to positive values
+                    similarities = similarities - similarities.min()
+                    # Normalize scores to sum to 100
+                    total = similarities.sum()
+                    if total > 0:
+                        normalized_scores = (similarities / total) * 100
+                    else:
+                        normalized_scores = similarities
+                    # Create a dictionary of filenames and scores
+                    scores = dict(zip(self.training_filenames, normalized_scores))
+                    # Sort the scores in descending order
+                    scores = dict(sorted(scores.items(), key=lambda item: item[1], reverse=True))
+                    # Limit the number of tracks if max_tracks is specified
+                    if max_tracks > 0:
+                        scores = dict(list(scores.items())[:max_tracks])
+                    similarity_scores.append(scores)
+
+            return similarity_scores 

training_filenames.txt

@@ -0,0 +1,10 @@
+ES_BUZZER - Cushy.wav
+ES_Life (Instrumental Version) - Northside.wav
+ES_A Christmas Dance - Arthur Benson.wav
+ES_Droppin Buckets (Instrumental Version) - Nyck Caution.wav
+ES_EPIC FIGHT SONG NO. 1 - Def Lev.wav
+ES_JACKHAMMER - Cushy.wav
+ES_Sunday Blues - Hara Noda.wav
+ES_Christmas Magic - Megan Wofford.wav
+ES_Moving Up - Origo.wav
+ES_Breeze - Basixx.wav