276 lines
8.6 KiB
Python
276 lines
8.6 KiB
Python
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()
|