| |
| """Test attribute conditioning: generate the SAME composer/key with CONTRASTING attribute sets |
| and check the output shifts in the requested direction (feature adherence). MPS, load once.""" |
| import os, sys, io, json |
| import torch, pretty_midi |
| ROOT = '/Users/alhill/projects/music/tmp/melody' |
| sys.path.insert(0, ROOT); sys.path.insert(0, f'{ROOT}/aria-repo') |
| from model_qwen import GPTQwen |
| from ariautils.tokenizer import AbsTokenizer |
| from attributes import extract, ATTR_ORDER |
|
|
| CK = torch.load(f'{ROOT}/ckpt_attr.pt', map_location='cpu'); A = CK['arch'] |
| META = json.load(open(f'{ROOT}/attr_meta.json')) |
| DEV = torch.device('mps' if torch.backends.mps.is_available() else 'cpu') |
| TOK = AbsTokenizer() |
| MODEL = GPTQwen(vocab=META['vocab'], d=A['d'], nh=A['nh'], nkv=A['nkv'], nl=A['nl'], ctx=A['ctx']).to(DEV) |
| MODEL.load_state_dict(CK['model']); MODEL.eval() |
| print(f"loaded ckpt_attr step {CK['step']} ppl {2.718281828**CK['val']:.2f} on {DEV}", flush=True) |
|
|
| INSTR, BOS, EOS = META['instr_prefix_id'], META['bos'], META['eos'] |
| COMP, KEYS, ATOK, AANY = META['composers'], META['keys'], META['attr_tokens'], META['attr_any'] |
| PREFIX = set(COMP.values()) | set(KEYS.values()) | {INSTR, BOS, EOS, META['pad_id']} |
| for at in ATTR_ORDER: |
| PREFIX |= set(ATOK[at].values()); PREFIX.add(AANY[at]) |
| MASK = torch.zeros(META['vocab']) |
| for i in PREFIX: |
| MASK[i] = float('-inf') |
|
|
|
|
| def attr_ids(attrs): |
| return [ATOK[at][attrs[at]] if at in attrs else AANY[at] for at in ATTR_ORDER] |
|
|
|
|
| def sample(logits, temp, topp): |
| lg = (logits.to('cpu').float() + MASK) / temp |
| p = torch.softmax(lg, dim=-1); sp, si = torch.sort(p, descending=True) |
| sp = sp * ((torch.cumsum(sp, 0) - sp) < topp) |
| return si[torch.multinomial(sp / sp.sum(), 1)].item() |
|
|
|
|
| @torch.no_grad() |
| def gen(composer, keymeta, attrs, n_tokens=400, temp=1.0, topp=0.96): |
| seed = [INSTR, COMP.get(composer, COMP['unknown']), KEYS[keymeta], BOS] + attr_ids(attrs) |
| toks = list(seed); notes = [] |
| logits, cache = MODEL.infer(torch.tensor([toks], device=DEV), cache=None, pos=0) |
| for _ in range(n_tokens): |
| nxt = sample(logits[0], temp, topp) |
| if nxt == EOS or len(toks) >= MODEL.ctx - 1: |
| break |
| toks.append(nxt) |
| if nxt not in PREFIX: |
| notes.append(nxt) |
| logits, cache = MODEL.infer(torch.tensor([[nxt]], device=DEV), cache=cache, pos=len(toks) - 1) |
| return notes |
|
|
|
|
| def to_notes(note_ids): |
| md = TOK.detokenize(TOK.decode([INSTR, BOS] + note_ids)); buf = io.BytesIO(); md.to_midi().save(file=buf) |
| pm = pretty_midi.PrettyMIDI(io.BytesIO(buf.getvalue())) |
| return sorted((n.start, n.pitch, n.end - n.start, n.velocity) for inst in pm.instruments for n in inst.notes) |
|
|
|
|
| TESTS = { |
| 'HIGH / DENSE / LOUD / ASCENDING': {'register': 'brilliant', 'density': 'torrential', 'dynamics': 'thunderous', |
| 'contour': 'ascending', 'motion': 'agitated', 'color': 'bright'}, |
| 'LOW / SPARSE / SOFT / DESCENDING': {'register': 'bass', 'density': 'sparse', 'dynamics': 'hushed', |
| 'contour': 'descending', 'motion': 'still', 'color': 'dark'}, |
| } |
| for name, attrs in TESTS.items(): |
| ni = gen('chopin', '0_maj', attrs) |
| notes = to_notes(ni) |
| got = extract(notes, spb=0.55) if len(notes) >= 3 else None |
| print(f'\n### {name} ({len(notes)} notes)') |
| if got: |
| for k in attrs: |
| hit = '✓' if got[k] == attrs[k] else ('~' if k in ('register', 'density', 'dynamics') else '✗') |
| print(f' {k:11} asked {attrs[k]:11} -> got {got[k]:11} {hit}') |
|
|
