| """Packen für M1/M2 mit dem SFT3-Prozessor (identisch zu sft_train2.build_examples, ohne Prompt-Zufall).""" |
| import sys, torch |
| sys.path.insert(0, '/e/data1/datasets/playground/mmlaion/schuhmann1/dramabox/train') |
| from moss_pack import collate |
| from prompt_fmt import user_content, AUDIO_PLACEHOLDER |
|
|
|
|
| class MossPacker: |
| def __init__(self, proc, cfg): |
| self.proc = proc |
| self.assist = int(cfg.audio_assistant_slot_token_id); self.aend = int(cfg.audio_end_token_id) |
| self.audio_pad = int(cfg.audio_pad_token_id); self.text_pad = int(cfg.pad_token_id) |
|
|
| def pack(self, meta, codes, ref_codes=None): |
| use_ref = ref_codes is not None |
| content = user_content(meta, use_ref) |
| um = {"role": "user", "content": content, "audio_codes_list": [torch.as_tensor(ref_codes)] if use_ref else []} |
| am = self.proc.build_assistant_message(audio_codes_list=[torch.as_tensor(codes)]) |
| b = self.proc([[um, am]], mode="computing_loss") |
| ids = b["input_ids"][0].long() |
| lab = torch.full_like(ids, -100) |
| lab[:-1] = ids[1:] |
| tt = lab[:, 0] |
| sup = (tt == self.assist) | (tt == self.aend) |
| fa = (ids[:, 0] == self.assist).nonzero() |
| if fa.numel() > 0: |
| sup[: max(0, int(fa[0]) - 1)] = False |
| lab[~sup] = -100 |
| au = lab[:, 1:] |
| au[(au >= self.audio_pad) | (au < 0)] = -100 |
| lab[:, 1:] = au |
| return {"input_ids": ids, "labels": lab, "n_sup": int(sup.sum())} |
|
|
| def prompt_ids(self, meta, ref_codes=None, frames=None): |
| """Generierungs-Eingabe [T,13] (endet mit <|audio_start|>).""" |
| use_ref = ref_codes is not None |
| content = user_content(meta, use_ref, frames=frames) |
| um = {"role": "user", "content": content, "audio_codes_list": [torch.as_tensor(ref_codes)] if use_ref else []} |
| b = self.proc([[um]], mode="generation") |
| return b["input_ids"][0].long() |
|
|
