# # # Batch GPU encode modal scripts # # import funcy import gc import os import pathlib import re import tempfile import time import torch import modal from modal import enter import numpy as np from scipy.special import softmax from suno_utils.utils.display import suppress_logging from contextlib import redirect_stderr from suno_utils.utils.s3 import read_from_s3, upload_s3_files, list_s3_dir from suno_utils.utils.text import read_jsonl, write_jsonl from suno_utils.tasks.data_loader import load_audio_mp from transformers import AutoModelForCausalLM, AutoTokenizer from huggingface_hub import InferenceClient N_MAX_REPLICAS = 1 image = ( modal.Image.debian_slim() .apt_install( "ffmpeg", "sox", "libsox-fmt-mp3", "libcairo2", "libcairo2-dev", "texlive-full" ) # .pip_install_private_repos( # "github.com/suno-ai/fairseq.git@9898715cdff75", # git_user="gkucsko", # secrets=[modal.Secret.from_name("glock_pat")], # ) # .pip_install_private_repos( # "github.com/suno-ai/hoot.git@1ad12a3", # git_user="gkucsko", # secrets=[modal.Secret.from_name("glock_pat")], # ) # .pip_install_from_pyproject( # str(pathlib.Path(__file__).parent.parent.parent.parent.parent / "descript-audio-codec/pyproject.toml"), # force_build=True, # ) .pip_install_from_pyproject( str(pathlib.Path(__file__).parent.parent.parent.parent / "pyproject.toml"), # force_build=True, ) .pip_install( "boto3", "transformers", "tokenizers", "encodec", "ctc_segmentation", "psutil", "nnAudio", "torch==2.1.1", "torchaudio==2.1.1", "importlib-resources==5.13.0", ) # .pip_install( # "stable-audio-tools", # ) # .pip_install_private_repos( # "github.com/suno-ai/glockenspiel.git@c35e74c58299#subdirectory=descript-audio-codec&egg=descript-audio-codec", # git_user="gkucsko", # secrets=[modal.Secret.from_name("glock90_pat")], # force_build=False, # ) ) stub = modal.App("batch-label-audio-corruptions", image=image) # ## Defining the prediction function # # * Container lifecycle hook: this lets us load the model only once in each container # # Instead of a using `@stub.function()` in the global scope, # we put the method on a class, and define an `__enter__` method on that class. # Modal reuses containers for successive calls to the same function, so # we want to take advantage of this and avoid setting up the same model # for every function call. # Notes: # relative speeds: 5.5s on a5000, 7.5s on A10G, 18.5s on T4 # n_cpu on T4/A10G: 24 # audio load: 75s for 1k files on 24 cores (70h audio) ~3-4k realtime @stub.cls( gpu="A10G", timeout=2 * 60 * 60, concurrency_limit=N_MAX_REPLICAS, secrets=[modal.Secret.from_name("aws-bucket")], # cloud="oci", memory=150000, # youtube 5k needs about 100gb retries=3, ) class Worker: def __init__( self, encode_s3_dir, ): self.encode_s3_dir = encode_s3_dir @enter() def set_up_models(self): @modal.method() def pseudo_label(self, work_item): work_item_id = work_item["id"] youtube_title = work_item["youtube_title"] # construct messages messages = [ { "role": "user", "content": f"{mistral_system_prompt} {mistral_example_title_1}", }, {"role": "assistant", "content": mistral_example_response_1}, {"role": "user", "content": f"{mistral_example_title_2}"}, {"role": "assistant", "content": mistral_example_response_2}, {"role": "user", "content": f"{mistral_example_title_3}"}, {"role": "assistant", "content": mistral_example_response_3}, {"role": "user", "content": f"{youtube_title}"}, ] with torch.no_grad(): encodeds = self.tokenizer.apply_chat_template(messages, return_tensors="pt") model_inputs = encodeds.to(self.device) generated_ids = self.model.generate( model_inputs, max_new_tokens=256, do_sample=True ) decoded = self.tokenizer.batch_decode(generated_ids) response = decoded[0].split("[/INST]")[-1] response = response.strip("") response = response.strip('"') response = response.strip("[") response = response.strip("]") tags = response.split(",") tags = [tag.strip() for tag in tags] tags = [tag.strip('"') for tag in tags] tags = [tag.strip(".") for tag in tags] new_tags = [] for tag in tags: # check for newlines if "\n" in tag: sub_tags = tag.split("\n") for sub_tag in sub_tags: new_tags.append(sub_tag) else: new_tags.append(tag) # save tags to a text file in comma separated format with open(f"{work_item_id}.txt", "w") as f: f.write(",".join(new_tags)) # upload the text file to the s3 bucket upload_s3_files(f"{work_item_id}.txt", self.encode_s3_dir) ## Run as eg: # modal run test_modal.py \ @stub.local_entrypoint() def main( s3_metas_filepath: str = "s3://suno-data/datasets/metadata/chirp_v4/genius_hq_metas_filtered.jsonl", output_name: str = "mistral_tags", chunksize: int = 500, first_only: bool = True, ): print("Downloading data...") metas = read_from_s3(s3_metas_filepath, read_f=read_jsonl) print(len(metas), "data items loaded") work_items = list(funcy.chunks(chunksize, metas)) work_items = list(zip(range(len(work_items)), work_items)) print(len(work_items), "work items") worker = Worker() print("Testing inference...") t0 = time.time() for work_item in work_items[:1]: _ = worker.pseudo_label.remote(work_item) print(f"{int(round(time.time()-t0))}s for test") if first_only: print("done with test.") return print("Running batch inference...") t0 = time.time() _ = list(worker.embed.map(work_items[1:])) print(round((time.time() - t0) / 60 / 60), "h total for batch embed") # verify that all is embedded done_archives = [fn for fn, _ in list_s3_dir(encode_s3_dir) if fn.endswith(".npz")] assert len(list(funcy.chunks(chunksize, range(len(metas))))) == len(done_archives) print(f"all {len(done_archives)} parts done.") # ## Further optimization notes # # Every container downloads the model when it starts, which is a bit inefficient. # In order to improve this, what you could do is to set up a shared volume that gets # mounted to each container. # See [shared volumes](/docs/guide/shared-volumes). # # In order for Huggingface to use the shared volume, you need to set the value of # the `TRANSFORMERS_CACHE` environment variable to the path of the shared volume. # See [secrets](/docs/guide/secrets).