#!/usr/bin/env python3 """ Detect Chinese songs in the training dataset using keywords and Chinese characters. Searches tags for Chinese-related keywords and text for Chinese characters. """ import argparse import json import re import time from datetime import datetime from pathlib import Path from typing import Dict, List, Set import random def setup_logging(output_dir: Path = None): """Set up basic logging.""" import logging if output_dir: output_dir.mkdir(parents=True, exist_ok=True) log_file = output_dir / f"chinese_detection_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log" logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s", handlers=[logging.FileHandler(log_file), logging.StreamHandler()], ) else: logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") return logging.getLogger(__name__) class ChineseSongDetector: """Detect Chinese songs in dataset.""" def __init__(self, dataset_file: str, output_dir: str = None, sample_rate: int = 1): """ Initialize detector. Args: dataset_file: Path to dataset JSONL output_dir: Output directory for results sample_rate: Sample 1 out of N records (1=all, 10=10%, 100=1%, etc.) """ self.dataset_file = Path(dataset_file) self.output_dir = Path(output_dir) if output_dir else Path("/home/vibert/tmp/chinese_detection") self.output_dir.mkdir(parents=True, exist_ok=True) self.sample_rate = sample_rate self.logger = setup_logging(self.output_dir) # Chinese-related keywords (case-insensitive) self.chinese_keywords = { # Language/dialect keywords "mandarin", "mando", "cantonese", "canto", "chinese", "taiwanese", "hokkien", "hakka", "putonghua", "普通话", # Country/region keywords "china", "taiwan", "hong kong", "hongkong", "beijing", "shanghai", "guangzhou", "shenzhen", "taipei", "macau", # Music genre keywords "c-pop", "cpop", "cantopop", "mandopop", "zhongguo feng", "chinese folk", "chinese traditional", "chinese classical", "chinese rock", "chinese rap", "chinese hip hop", # Cultural keywords "lunar new year", "spring festival", "mid-autumn", "dragon boat", "guzheng", "erhu", "pipa", "dizi", # Artist/label keywords (common ones) "jay chou", "teresa teng", "wang leehom", "jj lin", "g.e.m", "faye wong", "eason chan", "jacky cheung", } # Regex patterns for Chinese characters # CJK Unified Ideographs (main block): U+4E00-U+9FFF # CJK Extension A: U+3400-U+4DBF # Common punctuation: U+3000-U+303F self.chinese_char_pattern = re.compile(r"[\u4e00-\u9fff\u3400-\u4dbf]+") # Track statistics self.stats = { "total_records": 0, "records_processed": 0, "records_sampled": 0, "chinese_by_tags": 0, "chinese_by_text": 0, "chinese_by_both": 0, "chinese_total": 0, "tag_keyword_frequency": {}, "processing_time": 0, } # Store sample records for analysis self.chinese_samples = [] self.max_samples = 100 # Store up to 100 examples def has_chinese_keywords_in_tags(self, tags: List[str]) -> tuple[bool, List[str]]: """ Check if tags contain Chinese-related keywords. Returns: (has_chinese, matched_keywords) """ if not tags: return False, [] matched = [] tags_lower = " ".join(tags).lower() for keyword in self.chinese_keywords: if keyword in tags_lower: matched.append(keyword) # Track keyword frequency self.stats["tag_keyword_frequency"][keyword] = ( self.stats["tag_keyword_frequency"].get(keyword, 0) + 1 ) return len(matched) > 0, matched def has_chinese_characters_in_text(self, text: str) -> tuple[bool, int]: """ Check if text contains Chinese characters. Returns: (has_chinese, character_count) """ if not text: return False, 0 chinese_chars = self.chinese_char_pattern.findall(text) char_count = sum(len(chars) for chars in chinese_chars) # Consider it Chinese if more than 5 Chinese characters # (to avoid false positives from occasional characters) return char_count > 5, char_count def process_record(self, record: Dict) -> tuple[bool, Dict]: """ Process a single record to detect Chinese content. Returns: (is_chinese, detection_info) """ tags = record.get("tags", []) text = record.get("text", "") # Check tags has_chinese_tags, matched_keywords = self.has_chinese_keywords_in_tags(tags) # Check text for Chinese characters has_chinese_text, char_count = self.has_chinese_characters_in_text(text) is_chinese = has_chinese_tags or has_chinese_text detection_info = { "id": record.get("id", "unknown"), "has_chinese_tags": has_chinese_tags, "matched_keywords": matched_keywords, "has_chinese_text": has_chinese_text, "chinese_char_count": char_count, "is_chinese": is_chinese, } # Update statistics if has_chinese_tags: self.stats["chinese_by_tags"] += 1 if has_chinese_text: self.stats["chinese_by_text"] += 1 if has_chinese_tags and has_chinese_text: self.stats["chinese_by_both"] += 1 if is_chinese: self.stats["chinese_total"] += 1 # Store sample if under limit if len(self.chinese_samples) < self.max_samples: self.chinese_samples.append( { "id": record.get("id"), "detection": detection_info, "tags": tags[:10], # First 10 tags "text_preview": text[:200] if text else "", # First 200 chars } ) return is_chinese, detection_info def run(self): """Run detection on the dataset.""" self.logger.info("=" * 80) self.logger.info("CHINESE SONG DETECTION") self.logger.info("=" * 80) self.logger.info(f"Dataset: {self.dataset_file}") self.logger.info( f"Sample rate: 1/{self.sample_rate} (processing ~{100/self.sample_rate:.2f}% of data)" ) self.logger.info(f"Chinese keywords: {len(self.chinese_keywords)} keywords defined") start_time = time.time() # Process dataset self.logger.info("\nProcessing dataset...") with open(self.dataset_file, "r") as f: for line_num, line in enumerate(f, 1): self.stats["total_records"] += 1 # Apply sampling if self.sample_rate > 1: if line_num % self.sample_rate != 0: continue self.stats["records_sampled"] += 1 if not line.strip(): continue try: record = json.loads(line.strip()) self.stats["records_processed"] += 1 # Process record is_chinese, detection_info = self.process_record(record) # Progress update if self.stats["records_processed"] % 10000 == 0: elapsed = time.time() - start_time rate = self.stats["records_processed"] / elapsed chinese_pct = 100 * self.stats["chinese_total"] / self.stats["records_processed"] self.logger.info( f" Processed: {self.stats['records_processed']:,} | " f"Chinese: {self.stats['chinese_total']:,} ({chinese_pct:.2f}%) | " f"Rate: {rate:.0f} records/sec" ) except json.JSONDecodeError: continue except Exception as e: self.logger.warning(f"Error processing line {line_num}: {e}") self.stats["processing_time"] = time.time() - start_time # Generate report self.generate_report() def generate_report(self): """Generate and save detection report.""" # Calculate percentages if self.stats["records_processed"] > 0: chinese_pct = 100 * self.stats["chinese_total"] / self.stats["records_processed"] tags_only_pct = ( 100 * (self.stats["chinese_by_tags"] - self.stats["chinese_by_both"]) / self.stats["records_processed"] ) text_only_pct = ( 100 * (self.stats["chinese_by_text"] - self.stats["chinese_by_both"]) / self.stats["records_processed"] ) both_pct = 100 * self.stats["chinese_by_both"] / self.stats["records_processed"] else: chinese_pct = tags_only_pct = text_only_pct = both_pct = 0 # Estimate total if sampling was used if self.sample_rate > 1: estimated_total = self.stats["chinese_total"] * self.sample_rate estimated_pct = chinese_pct # Percentage should be the same else: estimated_total = self.stats["chinese_total"] estimated_pct = chinese_pct # Create report report = { "timestamp": datetime.now().isoformat(), "configuration": { "dataset_file": str(self.dataset_file), "sample_rate": self.sample_rate, "sample_percentage": 100 / self.sample_rate, "chinese_keywords_count": len(self.chinese_keywords), }, "statistics": { "total_records_in_file": self.stats["total_records"], "records_sampled": self.stats["records_sampled"], "records_processed": self.stats["records_processed"], "chinese_songs_found": self.stats["chinese_total"], "chinese_percentage": round(chinese_pct, 2), "chinese_by_tags_only": self.stats["chinese_by_tags"] - self.stats["chinese_by_both"], "chinese_by_text_only": self.stats["chinese_by_text"] - self.stats["chinese_by_both"], "chinese_by_both": self.stats["chinese_by_both"], "processing_time_seconds": round(self.stats["processing_time"], 2), "processing_rate": round( self.stats["records_processed"] / self.stats["processing_time"], 0 ), }, "estimates": { "estimated_total_chinese": estimated_total, "estimated_percentage": round(estimated_pct, 2), }, "top_keywords": dict( sorted(self.stats["tag_keyword_frequency"].items(), key=lambda x: x[1], reverse=True)[ :20 ] ), "sample_records": self.chinese_samples[:20], # First 20 samples } # Save JSON report report_file = self.output_dir / "chinese_detection_report.json" with open(report_file, "w", encoding="utf-8") as f: json.dump(report, f, indent=2, ensure_ascii=False) # Save text summary summary_file = self.output_dir / "chinese_detection_summary.txt" with open(summary_file, "w", encoding="utf-8") as f: f.write("=" * 80 + "\n") f.write("CHINESE SONG DETECTION RESULTS\n") f.write("=" * 80 + "\n\n") f.write("CONFIGURATION:\n") f.write(f" Dataset: {self.dataset_file.name}\n") f.write(f" Sample rate: 1/{self.sample_rate} ({100/self.sample_rate:.2f}% of data)\n") f.write(f" Processing time: {self.stats['processing_time']:.1f} seconds\n\n") f.write("RESULTS:\n") f.write(f" Total records in file: {self.stats['total_records']:,}\n") f.write(f" Records processed: {self.stats['records_processed']:,}\n") f.write(f" Chinese songs found: {self.stats['chinese_total']:,} ({chinese_pct:.2f}%)\n\n") f.write("DETECTION BREAKDOWN:\n") f.write( f" By tags only: {self.stats['chinese_by_tags'] - self.stats['chinese_by_both']:,} ({tags_only_pct:.2f}%)\n" ) f.write( f" By text only: {self.stats['chinese_by_text'] - self.stats['chinese_by_both']:,} ({text_only_pct:.2f}%)\n" ) f.write(f" By both: {self.stats['chinese_by_both']:,} ({both_pct:.2f}%)\n\n") if self.sample_rate > 1: f.write("ESTIMATES (based on sampling):\n") f.write(f" Estimated total Chinese songs: ~{estimated_total:,}\n") f.write(f" Estimated percentage: ~{estimated_pct:.2f}%\n\n") f.write("TOP MATCHED KEYWORDS:\n") for keyword, count in list(report["top_keywords"].items())[:10]: f.write(f" {keyword}: {count:,}\n") # Print summary self.logger.info("\n" + "=" * 80) self.logger.info("DETECTION COMPLETE") self.logger.info("=" * 80) self.logger.info(f"Records processed: {self.stats['records_processed']:,}") self.logger.info(f"Chinese songs found: {self.stats['chinese_total']:,} ({chinese_pct:.2f}%)") if self.sample_rate > 1: self.logger.info(f"Estimated total: ~{estimated_total:,} Chinese songs in full dataset") self.logger.info(f"Processing time: {self.stats['processing_time']:.1f} seconds") self.logger.info(f"\nReports saved to: {self.output_dir}") def main(): """Main execution.""" parser = argparse.ArgumentParser(description="Detect Chinese songs in training dataset") parser.add_argument( "--dataset-file", default="/app2/suno/data/auk_v0/metas_v6_tr.jsonl", help="Path to dataset JSONL", ) parser.add_argument( "--output-dir", default="/home/vibert/tmp/chinese_detection", help="Output directory for results" ) parser.add_argument( "--sample-rate", type=int, default=100, help="Sample 1 out of N records (1=all, 10=10%%, 100=1%%, 1000=0.1%%)", ) args = parser.parse_args() # Estimate time if args.sample_rate == 1: print("WARNING: Processing all ~50M records will take approximately 8-10 minutes") else: estimated_records = 50_000_000 / args.sample_rate estimated_time = estimated_records / 6000 # ~6000 records/sec print(f"Sampling 1/{args.sample_rate} records (~{estimated_records:,.0f} records)") print(f"Estimated time: {estimated_time:.1f} seconds") # Create detector detector = ChineseSongDetector( dataset_file=args.dataset_file, output_dir=args.output_dir, sample_rate=args.sample_rate ) # Run detection detector.run() if __name__ == "__main__": main()