import json import random import fire from csa import CSA import psycopg import uuid from hashlib import sha256 import typeguard from data import Vocab, VocabError, quantized_midi_to_indexable_bytes, wavtoolm2midi def dump_lanes(conn_str, vocab, train_split_idx, rank, size, out_path): with psycopg.connect(conn_str) as read_conn, psycopg.connect( conn_str ) as update_conn, read_conn.cursor(name=f"fetch_{uuid.uuid4()}") as cur, open( out_path, "wb" ) as outfile: update_conn.autocommit = True lane_hashes = [] def flush(toempty): if (not toempty and len(lane_hashes) < 1000) or len(lane_hashes) == 0: return with update_conn.cursor() as updcur, update_conn.transaction(): updcur.executemany( """ INSERT INTO lane_hashes (lane_id, hash) VALUES (%s, %b) ON CONFLICT (lane_id) DO UPDATE SET hash = EXCLUDED.hash WHERE lane_hashes.hash != EXCLUDED.hash """, lane_hashes, ) lane_hashes.clear() cur.itersize = 100 cur.execute( """ select l.id, f.split, jsonb_agg( jsonb_build_object( 'pitch', n.pitch, 'start', n.start, 'end', n."end", 'velocity', n.velocity ) ) from notes n join lanes l on l.id = n.lane_id join instruments i on i.id = l.instrument_id join files f on f.id = i.file_id where abs(hashint4extended(l.id, 12345)) %% %s = %s group by l.id, f.split """, (size, rank), ) for row in cur: try: qm = vocab.quantize_midi(wavtoolm2midi(row[2]), strict=False) except VocabError: continue if len(qm) == 0: continue qmb = quantized_midi_to_indexable_bytes(qm) lane_hashes.append((row[0], sha256(qmb).digest()[:16])) if row[1] == train_split_idx: outfile.write(qmb) flush(False) flush(True) def mark_lanes(conn_str, vocab, index, eval_split_idx, overlap_bytes=100): pending_lanes = [] count_processed = 0 count_marked_short = 0 count_marked_contaminated = 0 with psycopg.connect(conn_str) as read_conn, psycopg.connect( conn_str ) as update_conn, read_conn.cursor(name=f"fetch_{uuid.uuid4()}") as cur: update_conn.autocommit = True with update_conn.cursor() as updcur, update_conn.transaction(): updcur.execute( "INSERT INTO tag_values (value) VALUES ('Contaminated Eval') ON CONFLICT DO NOTHING" ) updcur.execute( "SELECT id FROM tag_values WHERE value = 'Contaminated Eval'" ) tag_value_id = updcur.fetchone()[0] cur.itersize = 100 cur.execute( """ select l.id, jsonb_agg( jsonb_build_object( 'pitch', n.pitch, 'start', n.start, 'end', n."end", 'velocity', n.velocity ) ) from notes n join lanes l on l.id = n.lane_id join instruments i on i.id = l.instrument_id join files f on f.id = i.file_id where f.split = %s group by l.id """, (eval_split_idx,), ) def flush(toempty): if (not toempty and len(pending_lanes) < 1000) or len(pending_lanes) == 0: return with update_conn.cursor() as updcur, update_conn.transaction(): updcur.executemany( """ INSERT INTO automatic_lane_tags (lane_id, tag_value_id) VALUES (%s, %s) ON CONFLICT DO NOTHING """, [(x, tag_value_id) for x in pending_lanes], ) pending_lanes.clear() for row in cur: count_processed += 1 flush(False) try: qm = vocab.quantize_midi(wavtoolm2midi(row[1]), strict=False) except VocabError: continue eval_data = quantized_midi_to_indexable_bytes(qm) if len(eval_data) < overlap_bytes: count_marked_short += 1 pending_lanes.append(row[0]) continue # per OpenAI - pick 3 random substrings and if any of them are in the training data, mark as contaminated for _ in range(3): start = random.randint(0, len(eval_data) - overlap_bytes) subbytes = eval_data[start : start + overlap_bytes] if index.count(subbytes) > 0: count_marked_contaminated += 1 pending_lanes.append(row[0]) break flush(True) print( f"Processed {count_processed} lanes, marked {count_marked_short} short and {count_marked_contaminated} contaminated" ) class DecontaminateTool: @typeguard.typechecked def dump_lanes( self, config_path: str, out_path: str, rank: int, size: int, ): """ Dumps a subset of lanes into indexable-byte-string format. Uses the configuration at config_path: - train_split_idx determines which files to index - quantize_divisions, pitch_{min,max}, drum_{min,max}, duration_max determine quantize behavior - dataset_class must be a DatasetFromPostgres*, and its conn_str will be used to connect to the database Output is written to out_path as concatenated indexable-byte-strings. Lane hashes are also written to the database for *all* lanes. """ assert size >= 1, "size must be at least 1" assert 0 <= rank and rank < size, "rank must be in [0, size)" with open(config_path, "r") as f: config = json.load(f) vocab = Vocab.from_config(config) conn_str = config["dataset_class"].get("conn_str", None) if conn_str is None: raise ValueError("conn_str not found in dataset_class") train_split_idx = config["train_split_idx"] dump_lanes(conn_str, vocab, train_split_idx, rank, size, out_path) @typeguard.typechecked def build_index( self, in_path: str, out_path: str, ): """ Builds an index from a concatenated set of indexable-byte-strings. in_path is the path to the concatenated indexable-byte-strings. Output is written to out_path as a serialized sdsl csa index. """ CSA.build(in_path, out_path) @typeguard.typechecked def mark_lanes( self, config_path: str, index_path: str, ): """ Marks all lanes in the eval split that are contaminated. Uses the configuration at config_path: - eval_split_idx determines which files to mark - quantize_divisions, pitch_{min,max}, drum_{min,max}, duration_max determine quantize behavior - dataset_class must be a DatasetFromPostgres, and its conn_str will be used to connect to the database index_path is the path to the serialized sdsl csa index built from the training split. A "Contaminated Eval" tag will be added to automatic_lane_tags for each contaminated lane. """ with open(config_path, "r") as f: config = json.load(f) vocab = Vocab.from_config(config) conn_str = config["dataset_class"].get("conn_str", None) if conn_str is None: raise ValueError("conn_str not found in dataset_class") eval_split_idx = config["eval_split_idx"] index = CSA.load(index_path) mark_lanes(conn_str, vocab, index, eval_split_idx) if __name__ == "__main__": from util import configure_logging configure_logging() fire.Fire(DecontaminateTool())