from concurrent.futures import ProcessPoolExecutor, as_completed from tqdm import tqdm from datasketch import MinHash, MinHashLSH from collections import defaultdict import random import re import unicodedata import hashlib from datasketch import MinHash, MinHashLSH from collections import defaultdict import os import json from joblib import Parallel, delayed from suno_utils.utils.text import read_jsonl # ------------------------------- # 1. Normalization # ------------------------------- def normalize(text): # Unicode normalize + casefold text = unicodedata.normalize("NFKD", text).casefold() text = "".join(c for c in text if not unicodedata.combining(c)) # Remove bracketed content like "(Official Video)", "[Remix]" text = re.sub(r"\([^)]*\)|\[[^]]*\]", " ", text) # REMOVE common noise words early text = re.sub( r"\b(remaster(ed)?|live|mono|stereo|version|edit|single|official|video|hd|feat\.?|intro|outro|instrumental|radio|mix|remix)\b", " ", text, ) # Split on common separators and normalize each part parts = re.split(r"[-–—|]", text) parts = [p.strip() for p in parts if p.strip()] if len(parts) == 2: # Create both possible orderings return [ set(parts[0].split() + parts[1].split()), # all words together set(parts[1].split() + parts[0].split()), # reversed order ] else: # If not exactly 2 parts, just return the words return [set(text.split())] def block_key(text, size=6): # If text is a list of sets, use the first set's words joined together if isinstance(text, list) and isinstance(text[0], set): text = " ".join(text[0]) # Split on common separators parts = re.split(r"[-–—|]", text) parts = [p.strip() for p in parts if p.strip()] if len(parts) == 2: # Get consonants from both parts cons1 = re.sub(r"[aeiou\s]", "", parts[0]) cons2 = re.sub(r"[aeiou\s]", "", parts[1]) # Take first few consonants from each part key1 = cons1[: size // 2] key2 = cons2[: size // 2] # Sort the keys to make it order-independent keys = sorted([key1, key2]) # Combine and hash combined = "".join(keys) else: # If not exactly 2 parts, just use the first few consonants consonants = re.sub(r"[aeiou\s]", "", text) combined = consonants[:size] return hashlib.md5(combined.encode()).hexdigest()[:8] class UnionFind: def __init__(self): self.parent = dict() def find(self, x): if x not in self.parent: self.parent[x] = x if self.parent[x] != x: self.parent[x] = self.find(self.parent[x]) return self.parent[x] def union(self, x, y): self.parent[self.find(x)] = self.find(y) # Reuse these from before def char_ngrams(s, n=3): s = f" {s} " return {s[i : i + n] for i in range(len(s) - n + 1)} def create_minhash(ngrams, num_perm=128): m = MinHash(num_perm=num_perm) for g in ngrams: m.update(g.encode("utf8")) return m def process_block(block, jaccard_thresh=0.92): pairs = set() if len(block) <= 1: return pairs # Create word sets for each title word_sets = {} for id_, text in block: word_sets[id_] = normalize(text) # Compare all pairs for id1 in word_sets: for id2 in word_sets: if id1 >= id2: continue # Try all combinations of normalizations max_sim = 0 for words1 in word_sets[id1]: for words2 in word_sets[id2]: # Skip if either set is empty if not words1 or not words2: continue sim = len(words1 & words2) / len(words1 | words2) max_sim = max(max_sim, sim) if max_sim >= jaccard_thresh: pairs.add((id1, id2)) return pairs def cluster_records_parallel(data, jaccard_threshold=0.92, lsh_threshold=0.95): norm_map = {d["id"]: normalize(d["title"]) for d in data} blocks = defaultdict(list) for id_, text in norm_map.items(): blocks[block_key(text)].append((id_, text)) uf = UnionFind() # Parallel block processing all_pairs = [] with ProcessPoolExecutor() as executor: futures = { executor.submit(process_block, block, jaccard_threshold): block for block in blocks.values() } for future in tqdm( as_completed(futures), total=len(futures), desc="Clustering blocks" ): result = future.result() all_pairs.extend(result) # Union-Find for id1, id2 in all_pairs: uf.union(id1, id2) # Collect clusters clusters = defaultdict(list) for id_ in norm_map: clusters[uf.find(id_)].append(id_) return [sorted(cluster) for cluster in clusters.values() if len(cluster) > 1] def cluster_records_joblib(data, jaccard_threshold=0.92, block_size=10, n_jobs=-1): # Create normalized versions for each title norm_map = {} for d in data: norm_map[d["id"]] = normalize(d["title"]) # Build blocks using the first normalized version for blocking blocks = defaultdict(list) for id_, norm_versions in tqdm(norm_map.items(), desc="Building blocks"): # Use the first normalized version for blocking first_norm = " ".join(norm_versions[0]) blocks[block_key(first_norm, block_size)].append((id_, first_norm)) uf = UnionFind() block_list = list(blocks.values()) # tqdm-compatible joblib print("Processing blocks in parallel...") results = Parallel(n_jobs=n_jobs, prefer="processes")( delayed(process_block)(block, jaccard_threshold) for block in tqdm(block_list, desc="Clustering blocks") ) # Flatten pairs and build unions for pairs in results: for id1, id2 in pairs: uf.union(id1, id2) # Collect clusters clusters = defaultdict(list) for id_ in norm_map: clusters[uf.find(id_)].append(id_) return [sorted(cluster) for cluster in clusters.values() if len(cluster) > 1] if __name__ == "__main__": # now we load the original metas to get the artists and song titles # METAS_DIR = "/app/suno/tmp" # raw_metas = read_jsonl(os.path.join(METAS_DIR, "raw_discogs_subset_metas.jsonl")) # raw_metas = read_jsonl(os.path.join(METAS_DIR, "raw_genius_metas.jsonl")) # raw_metas = read_jsonl(os.path.join(METAS_DIR, "raw_imslp_metas.jsonl")) # print(len(raw_metas)) # dataset_name = "imslp" dataset = "combined" # dataset = "youtube_music" out_dir = "/home/christian/code/christian/metadata/dedup" if not os.path.exists(out_dir): os.makedirs(out_dir) # get the titles title_filepaths = [ "/home/christian/code/christian/metadata/dedup/imslp_titles.json", "/home/christian/code/christian/metadata/dedup/deezer_titles.json", "/home/christian/code/christian/metadata/dedup/youtube_music_titles.json", "/home/christian/code/christian/metadata/dedup/pond5_titles.json", "/home/christian/code/christian/metadata/dedup/discogs_titles.json", "/home/christian/code/christian/metadata/dedup/genius_titles.json", "/home/christian/code/christian/metadata/dedup/discogs_subset_titles.json", ] all_titles = [] for title_filepath in title_filepaths: with open(title_filepath, "r") as f: titles = json.load(f) print(f"Loaded {len(titles)} titles from {title_filepath}") all_titles.extend(titles) # test on a subset first # all_titles = random.sample(all_titles, 1_000_000) print(f"Total titles: {len(all_titles)}") id2title = {d["id"]: d["title"] for d in all_titles} # cluster the titles jaccard_threshold = 0.80 block_size = 10 clusters = cluster_records_joblib( all_titles, jaccard_threshold=jaccard_threshold, block_size=block_size, n_jobs=-1, ) print(f"Found {len(clusters)} clusters") # Prepare data for JSON serialization clusters_with_titles = [] for cluster in clusters: cluster_data = { "size": len(cluster), "items": [{"id": id_, "title": id2title[id_]} for id_ in cluster], } clusters_with_titles.append(cluster_data) # Save clusters to a JSON file output_file = os.path.join( out_dir, f"{dataset}_title_clusters_jac-{jaccard_threshold:.2f}.json", ) with open(output_file, "w") as f: json.dump(clusters_with_titles, f, indent=2) print(f"Saved {len(clusters)} clusters to {output_file}")