"""THIS FILE JUST DOES ALIGNMENT...FAST Takes the saved words times and aligns them. This is flexible as we can adjust the algiment's max duration. NOTE That this file automatically join and merges the final jsonl files. """ import os import json import argparse import tqdm from tqdm.contrib.concurrent import process_map from multiprocessing import cpu_count import random import time from suno_utils.tasks.hoot import get_aligned_lyrics from suno_utils.utils.text import write_jsonl, read_jsonl from suno_utils.utils.tokenizers import tokenize import numpy as np import shutil def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--dataset", type=str) args = parser.parse_args() return args if __name__ == "__main__": input_args = parse_args() version_number = "h5_t480_v5" # For fixed chunk data preparation # use _collect_valid_segment_from_lines_by_duration in hoot.py TMP_DIR = "/home/tony/Work/tony/hoot/tmp" if input_args.dataset == "genius": INPUT_PATH = "/home/tony/Data/Hoot/genius_metas.json" OUTPUT_PATH = "/app/suno/data/hoot/alignments/genius_v5_fix" ALIGNED_OUTPUT_PATH = f"/home/tony/Data/Hoot/alignments/genius_{version_number}" elif input_args.dataset == "ytm": INPUT_PATH = "/home/tony/Data/Hoot/ytm_metas.jsonl" OUTPUT_PATH = "/app/suno/data/hoot/ytm_3" ALIGNED_OUTPUT_PATH = f"/home/tony/Data/Hoot/alignments/ytm_{version_number}" elif input_args.dataset == "deezer": INPUT_PATH = "/home/tony/Data/Hoot/deezer_metas.json" OUTPUT_PATH = "/app/suno/data/hoot/alignments/deezer_v5_fix" ALIGNED_OUTPUT_PATH = f"/home/tony/Data/Hoot/alignments/deezer_{version_number}" elif input_args.dataset == "discogs": INPUT_PATH = "/home/tony/Data/Hoot/discogs_full_metas.json" OUTPUT_PATH = "/app/suno/data/hoot/alignments/discogs_full_v5_fix" ALIGNED_OUTPUT_PATH = ( f"/home/tony/Data/Hoot/alignments/discogs_{version_number}" ) else: raise ValueError("Unknown dataset") # load the infomration with open(INPUT_PATH, "r") as fp: metas = json.load(fp) meta_id_map = {meta["id"]: meta for meta in metas} # with open("/home/tony/Data/Hoot/alignments/all_end_time.json", "r") as fp: # all_end_time = json.load(fp) # Remove ALIGNED_OUTPUT_PATH directory and all contents if it exists # if os.path.exists(ALIGNED_OUTPUT_PATH): # shutil.rmtree(ALIGNED_OUTPUT_PATH) # make the aligned output path os.makedirs(ALIGNED_OUTPUT_PATH, exist_ok=True) def generatge_alinged_lyrics_map(index): try: random.seed(index) output_jsonl_path = f"{ALIGNED_OUTPUT_PATH}/batch_{index}.jsonl" if os.path.exists(output_jsonl_path): # let's avoid the race condition time.sleep(0.1) return # print("index") with open(output_jsonl_path, "w") as fp: fp.write("") time.sleep(0.1) # this is just doing alignment if not os.path.exists(f"{OUTPUT_PATH}/batch_{index}.json"): return with open(f"{OUTPUT_PATH}/batch_{index}.json", "r") as fp: hooted_outputs = json.load(fp) # # if you want to refine the endtime # refined_hoot_outputs = [] local_id_to_cer_map = {} for meta in hooted_outputs: meta_id = meta["id"] # up to v3 we still need this hack..as the original lyrics are cased meta["lyrics"] = ( meta_id_map[meta_id]["lyrics"] if "lyrics" in meta_id_map[meta_id] else meta_id_map[meta_id]["text"] ) local_id_to_cer_map[meta_id] = meta["cer_val"] if meta_id == "W3JRaGV42Rg": print("FOUND", meta) # print(local_id_to_cer_map) # print(len(hooted_outputs))q # filter some inputs aligned_lyrics_map = get_aligned_lyrics( hooted_outputs, silent=True, min_cer=0.99, min_p_align=0.0, aligned_text_frac=0.99, return_only_valid=True, ) # print(len(hooted_outputs), len(aligned_lyrics_map)) write_jsonl( [(k, v, local_id_to_cer_map[k]) for k, v in aligned_lyrics_map.items()], output_jsonl_path, ) except Exception as e: print(f"WTF {index}, {e}") time.sleep(0.1) print("Start aligning lyrics.") input_jobs = [] total_files = len(os.listdir(OUTPUT_PATH)) for i in tqdm.tqdm(range(0, total_files)): # this is for testing only # for i in tqdm.tqdm(range(0, 100, 100)): input_jobs.append(i * 100) # break print(f"Total jobs: {len(input_jobs)}") # for debug # input_jobs = input_jobs[:10] # print(input_jobs) process_map( generatge_alinged_lyrics_map, input_jobs, max_workers=cpu_count() - 4, chunksize=1, ) print("DONE!!!") dataset_name = input_args.dataset aligned_output_folder = ( f"/home/tony/Data/Hoot/alignments/{dataset_name}_{version_number}" ) output_large_jsonl_name = f"{dataset_name}_hq_alignments_{version_number}.jsonl" for n_step, jsonl_name in tqdm.tqdm( enumerate(sorted(os.listdir(aligned_output_folder))) ): try: flat_aligned_lyrics = read_jsonl( os.path.join(aligned_output_folder, jsonl_name) ) write_jsonl( flat_aligned_lyrics, os.path.join(TMP_DIR, output_large_jsonl_name), do_append=bool(n_step != 0), ) except Exception as e: print(e) print(n_step, jsonl_name) # raise ValueError("WTF") print(f"DONE!!! {os.path.join(TMP_DIR, output_large_jsonl_name)}") # get summary (sample) stats with open(os.path.join(TMP_DIR, output_large_jsonl_name)) as f: n_tot_rows = sum(1 for _ in f) test_tokenize = True n_max_sample = n_tot_rows if test_tokenize: n_max_sample = int(round(n_max_sample / 180)) sample_data = [] tot_duration_s = 0 tot_duration_silence_s = 0 n_tokens_list = [] durations_list = [] unique_ids = set() n = 0 with open(os.path.join(TMP_DIR, output_large_jsonl_name)) as f: for line in tqdm.tqdm(f, total=n_tot_rows): line = line.strip() if len(line) == 0: continue k, v, cer_val = json.loads(line) if k in unique_ids: continue unique_ids.add(k) for e in v: if e["text"] != "": tot_duration_s += e["end_s"] - e["start_s"] durations_list.append(e["end_s"] - e["start_s"]) if e["text"] == "": tot_duration_silence_s += e["end_s"] - e["start_s"] if test_tokenize: n_tokens_list.append(len(tokenize(e["text"], max_tokens=512 * 8))) if n % 500 == 0: sample_data.append((k, v)) n += 1 if n == n_max_sample: break # print(int(len(unique_ids)*n_tot_rows/n_max_sample), "items") print(round(tot_duration_s * n_tot_rows / n_max_sample / 60 / 60, 1), "hours") print( round(tot_duration_silence_s * n_tot_rows / n_max_sample / 60 / 60, 1), "hours silence", ) print( f"tokens median {np.median(n_tokens_list):.1f}, 99% {np.quantile(n_tokens_list, 0.99):.1f}" ) print( f"durations median {np.median(durations_list):.1f}, 90% {np.quantile(durations_list, 0.90):.1f}" )