初始化项目，由ModelHub XC社区提供模型

Model: jy1095/qwen3-0.6b-neucodec-multipack-test
Source: Original Platform

qwen_train_subset_multipack.py

@@ -0,0 +1,275 @@
+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"<speech_{code}>" for code in codes)
+    prompt = f"<speech_start> {speech_text} <speech_end>\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"<speech_{i}>" for i in range(NUM_SPEECH_TOKENS)]
+    tokenizer.add_tokens(["<speech_start>", "<speech_end>"] + 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()