import csv import io import zipfile from pathlib import Path import pandas as pd import torch from torch.optim import AdamW from transformers import AutoModelForCausalLM, AutoTokenizer DATASET_DIR = Path("/workspace/fleurs-r-neucodec") MODEL_NAME = "Qwen/Qwen3-0.6B" NUM_SPEECH_TOKENS = 65536 MAX_SPEECH_TOKENS = 500 MAX_LENGTH = 1536 TRAIN_SPLIT = "train" VAL_SPLIT = "dev" MAX_TRAIN_EXAMPLES = 500 MAX_VAL_EXAMPLES = 50 LR = 1e-5 EPOCHS = 1 EVAL_EVERY = 25 SAVE_DIR = Path("/workspace/qwen_speech_multipack_2_ckpt") LOG_CSV = Path("/workspace/train_log_multipack_2.csv") def list_token_zips(split): zips = sorted((DATASET_DIR / "neucodec").glob(f"en_us-{split}*.zip")) if not zips: raise FileNotFoundError(f"No token zips found for split={split}") return zips def build_zip_index(zip_paths): index = {} open_zips = [] for path in zip_paths: zf = zipfile.ZipFile(path) open_zips.append(zf) for name in zf.namelist(): if name.endswith(".pt"): stem = Path(name).stem index[stem] = (zf, name) return index, open_zips def load_codes(zip_index, neucodec_path): stem = Path(str(neucodec_path).replace("\\", "/")).stem if stem not in zip_index: raise FileNotFoundError(f"No token file found for {neucodec_path}") zf, entry = zip_index[stem] obj = torch.load(io.BytesIO(zf.read(entry)), map_location="cpu") return obj["codes"].flatten().to(torch.long).tolist() def build_single_example(tokenizer, codes, transcript): codes = codes[:MAX_SPEECH_TOKENS] speech_text = " ".join(f"" for code in codes) prompt = f" {speech_text} \n" target = str(transcript) + tokenizer.eos_token prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"] target_ids = tokenizer(target, add_special_tokens=False)["input_ids"] input_ids = prompt_ids + target_ids labels = [-100] * len(prompt_ids) + target_ids input_ids = input_ids[:MAX_LENGTH] labels = labels[:MAX_LENGTH] return input_ids, labels def pack_examples(single_examples): packed = [] cur_input_ids = [] cur_labels = [] cur_segment_ids = [] segment_id = 0 for input_ids, labels in single_examples: if not input_ids: continue if cur_input_ids and len(cur_input_ids) + len(input_ids) > MAX_LENGTH: packed.append( { "input_ids": torch.tensor(cur_input_ids, dtype=torch.long), "labels": torch.tensor(cur_labels, dtype=torch.long), "segment_ids": torch.tensor(cur_segment_ids, dtype=torch.long), } ) cur_input_ids = [] cur_labels = [] cur_segment_ids = [] segment_id = 0 if len(input_ids) > MAX_LENGTH: input_ids = input_ids[:MAX_LENGTH] labels = labels[:MAX_LENGTH] cur_input_ids.extend(input_ids) cur_labels.extend(labels) cur_segment_ids.extend([segment_id] * len(input_ids)) segment_id += 1 if cur_input_ids: packed.append( { "input_ids": torch.tensor(cur_input_ids, dtype=torch.long), "labels": torch.tensor(cur_labels, dtype=torch.long), "segment_ids": torch.tensor(cur_segment_ids, dtype=torch.long), } ) return packed def load_examples(tokenizer, split, max_examples): parquet = DATASET_DIR / "data" / f"en_us-{split}.parquet" df = pd.read_parquet(parquet).head(max_examples) zip_paths = list_token_zips(split) zip_index, open_zips = build_zip_index(zip_paths) single_examples = [] for _, row in df.iterrows(): codes = load_codes(zip_index, row["neucodec_path"]) single_examples.append(build_single_example(tokenizer, codes, row["sentence"])) packed_examples = pack_examples(single_examples) return packed_examples, open_zips def make_block_causal_mask(segment_ids, dtype): # segment_ids: [L]. Tokens can attend only to earlier tokens in the same packed example. segment_ids = segment_ids.cuda() length = segment_ids.numel() same_segment = segment_ids[:, None] == segment_ids[None, :] causal = torch.arange(length, device="cuda")[:, None] >= torch.arange(length, device="cuda")[None, :] allowed = same_segment & causal mask = torch.zeros((1, 1, length, length), device="cuda", dtype=dtype) mask = mask.masked_fill(~allowed[None, None, :, :], torch.finfo(dtype).min) return mask def make_position_ids(segment_ids): # Reset positions at each packed-example boundary. position_ids = torch.zeros_like(segment_ids) for segment in torch.unique(segment_ids): idx = torch.nonzero(segment_ids == segment, as_tuple=False).flatten() position_ids[idx] = torch.arange(idx.numel(), dtype=torch.long) return position_ids.unsqueeze(0).cuda() @torch.inference_mode() def evaluate(model, examples): model.eval() losses = [] for ex in examples: input_ids = ex["input_ids"].unsqueeze(0).cuda() labels = ex["labels"].unsqueeze(0).cuda() attention_mask = make_block_causal_mask(ex["segment_ids"], model.dtype) position_ids = make_position_ids(ex["segment_ids"]) out = model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, labels=labels, ) losses.append(float(out.loss)) model.train() return sum(losses) / len(losses) def main(): tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) tokenizer.pad_token = tokenizer.eos_token speech_tokens = [f"" for i in range(NUM_SPEECH_TOKENS)] tokenizer.add_tokens(["", ""] + speech_tokens) print("Loading and multipacking train examples...") train_examples, train_zips = load_examples(tokenizer, TRAIN_SPLIT, MAX_TRAIN_EXAMPLES) print("Loading and multipacking validation examples...") val_examples, val_zips = load_examples(tokenizer, VAL_SPLIT, MAX_VAL_EXAMPLES) print(f"packed train batches: {len(train_examples)}") print(f"packed val batches: {len(val_examples)}") print(f"vocab size: {len(tokenizer)}") model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, torch_dtype=torch.bfloat16, trust_remote_code=True, ) model.resize_token_embeddings(len(tokenizer)) model.cuda() model.train() optimizer = AdamW(model.parameters(), lr=LR) with LOG_CSV.open("w", newline="") as f: writer = csv.DictWriter(f, fieldnames=["step", "train_loss", "val_loss"]) writer.writeheader() step = 0 for epoch in range(EPOCHS): for ex in train_examples: step += 1 input_ids = ex["input_ids"].unsqueeze(0).cuda() labels = ex["labels"].unsqueeze(0).cuda() attention_mask = make_block_causal_mask(ex["segment_ids"], model.dtype) position_ids = make_position_ids(ex["segment_ids"]) out = model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, labels=labels, ) loss = out.loss loss.backward() optimizer.step() optimizer.zero_grad(set_to_none=True) train_loss = float(loss.detach()) val_loss = "" if step % EVAL_EVERY == 0: val_loss = evaluate(model, val_examples) print(f"step {step:04d} train_loss {train_loss:.4f} val_loss {val_loss:.4f}") else: print(f"step {step:04d} train_loss {train_loss:.4f}") with LOG_CSV.open("a", newline="") as f: writer = csv.DictWriter(f, fieldnames=["step", "train_loss", "val_loss"]) writer.writerow( { "step": step, "train_loss": train_loss, "val_loss": val_loss, } ) SAVE_DIR.mkdir(parents=True, exist_ok=True) model.save_pretrained(SAVE_DIR) tokenizer.save_pretrained(SAVE_DIR) print(f"saved checkpoint: {SAVE_DIR}") print(f"saved log: {LOG_CSV}") for zf in train_zips + val_zips: zf.close() if __name__ == "__main__": main()