Buckets:
| """Many Veena takes for a few hard lines; keeps every take (+ both transcripts) so the final pick can be re-checked locally.""" | |
| import os, sys, json, tarfile | |
| import numpy as np, torch, soundfile as sf | |
| sys.path.insert(0, '/data/code'); import eval_lib as E | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| from snac import SNAC | |
| SPK = os.environ.get('SPK', 'vinaya'); N = int(os.environ.get('TAKES', '8')); R = int(os.environ.get('ROUNDS', '3')) | |
| TAG = os.environ.get('TAG', 'many'); LINES = json.load(open(os.environ['LINES'])) | |
| TEMPS = [float(x) for x in os.environ.get('TEMPS', '0.4').split(',')] | |
| SR = 24000; BASE = 128266 | |
| snac = SNAC.from_pretrained('hubertsiuzdak/snac_24khz').eval().to('cuda') | |
| tok = AutoTokenizer.from_pretrained('maya-research/Veena') | |
| m = AutoModelForCausalLM.from_pretrained('maya-research/Veena', torch_dtype=torch.bfloat16, device_map='cuda').eval() | |
| def decode(t): | |
| t = [x for x in t if BASE <= x < BASE + 7 * 4096]; n = len(t) // 7 | |
| if n < 4: return None | |
| l1, l2, l3 = [], [], [] | |
| for i in range(n): | |
| s = [t[i * 7 + j] - BASE - j * 4096 for j in range(7)] | |
| if any(x < 0 or x > 4095 for x in s): s = [(t[i * 7 + j] - BASE) % 4096 for j in range(7)] | |
| l1.append(s[0]); l2 += [s[1], s[4]]; l3 += [s[2], s[3], s[5], s[6]] | |
| codes = [torch.tensor(x, dtype=torch.long, device='cuda')[None] for x in (l1, l2, l3)] | |
| with torch.inference_mode(): a = snac.decode(codes) | |
| return a.squeeze().float().clamp(-1, 1).cpu().numpy() | |
| def gen(text, n, temp): | |
| ids = torch.tensor([[128259] + tok.encode(f'<spk_{SPK}> {text}', add_special_tokens=False) + [128260, 128261, 128257]], device='cuda') | |
| with torch.inference_mode(): | |
| o = m.generate(ids, attention_mask=torch.ones_like(ids), do_sample=True, num_return_sequences=n, eos_token_id=[128258, 128262], pad_token_id=128258, | |
| max_new_tokens=900, temperature=temp, top_p=0.9, repetition_penalty=1.05) | |
| outs = [] | |
| for i in range(n): | |
| s = o[i, ids.shape[1]:].tolist(); cut = next((j for j, x in enumerate(s) if x in (128258, 128262)), len(s)); outs.append(decode(s[:cut])) | |
| return outs | |
| d = f'/data/out/{SPK}_{TAG}'; os.makedirs(f'{d}/takes', exist_ok=True); rep = {} | |
| for L in LINES: | |
| k, ref, mx = L['id'], L['ref'], L['max']; rows = []; c = 0 | |
| for r in range(R): | |
| for vi, vt in enumerate(L['tts']): | |
| for temp in TEMPS: | |
| for a in gen(vt, N, temp): | |
| if a is None or len(a) < SR * .3: continue | |
| a = E.trim(a, SR); ev = E.evaluate(a, SR, ref); c += 1 | |
| fn = f'{k}_{c:03d}.flac'; sf.write(f'{d}/takes/{fn}', a, SR) | |
| rows.append(dict(file=fn, variant=vi, temp=temp, **{x: ev[x] for x in ('parakeet', 'moonshine', 'wer_p', 'wer_m', 'indian', 'accent', 'dur')})) | |
| print(f" {k} {fn} v{vi} t{temp} P={ev['wer_p']:.2f} M={ev['wer_m']:.2f} dur={ev['dur']:.2f} | {ev['parakeet']} | {ev['moonshine']}", flush=True) | |
| if sum(1 for x in rows if x['wer_p'] == 0 and x['wer_m'] == 0 and x['dur'] <= mx) >= int(os.environ.get('NEED', '4')): break | |
| rep[k] = rows; print(k, 'passes', sum(1 for x in rows if x['wer_p'] == 0 and x['wer_m'] == 0), 'of', len(rows), flush=True) | |
| json.dump(rep, open(f'{d}/report.json', 'w'), indent=1, ensure_ascii=False) | |
| with tarfile.open(f'/pub/{SPK}_{TAG}.tar.gz', 'w:gz') as tf: tf.add(d, arcname=f'{SPK}_{TAG}') | |
| print('ALL DONE') | |
Xet Storage Details
- Size:
- 3.44 kB
- Xet hash:
- 488e815a219913cfd8486c2119fc96fc2550f82b08588ccec473f2cabc21eb4b
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.