FilesExpand file tree

main

/

check_moss_codec_contexts.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

67 lines (60 loc) · 3.05 KB

main

/

check_moss_codec_contexts.py

Copy path

Top

File metadata and controls

• Code
• Blame

67 lines (60 loc) · 3.05 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

"""Check the MOSS codec attention contexts against the checkpoints they describe.

audio.cpp hardcodes each codec's per-stage attention context as a step count.

The checkpoint states it as a DURATION, per stage, which the reference turns into

steps with `int(round(current_frame_rate * context_duration))` -- where the rate

is the one *before* that transformer, not after its patch.

Getting that wrong is silent: the decoder still runs, still produces speech, and

only diverges once a sequence outgrows the real window, so short clips look

perfect and long ones drift. That is exactly what #663 reported, and it was the

Nano decoder, whose contexts were each twice what they should have been.

Needs the checkpoints' config.json files, so it is a tool rather than a test:

python3 tools/check_moss_codec_contexts.py

"""

import json, re, sys

def derive(cfg_path):

c = json.load(open(cfg_path))

default = float(c.get("causal_transformer_context_duration", 10.0))

ci = 2 if c.get("enable_channel_interleave") else 1

code_rate = c["sampling_rate"] / c["downsample_rate"]

out = {}

# decoder: starts at the code rate, rate *= patch after each module

rate = code_rate; dec = []

for m in c["decoder_kwargs"]:

if m["module_type"] == "PatchedPretransform":

rate *= m["patch_size"]

else:

dec.append(int(round(rate * float(m.get("context_duration", default)))))

out["decoder"] = (dec, rate, c["sampling_rate"] * ci)

# encoder: starts at the output rate, rate /= patch

rate = c["sampling_rate"] * ci; enc = []

for m in c["encoder_kwargs"]:

if m["module_type"] == "PatchedPretransform":

rate /= m["patch_size"]

else:

enc.append(int(round(rate * float(m.get("context_duration", default)))))

out["encoder"] = (enc, rate, code_rate)

return out

def hardcoded(src, fn):

body = src.split(f"MossAudioTokenizerConfig {fn}()")[1].split("\n}")[0]

out = {}

for side in ("encoder_stages", "decoder_stages"):

block = body.split(f"config.{side} = {{")[1].split("};")[0]

out[side.split("_")[0]] = [int(r.split(",")[6]) for r in

re.findall(r"\{([^}]*)\}", block)]

return out

src = open("src/framework/codecs/moss_audio_tokenizer_codec_runtime.cpp").read()

for label, cfg, fn in [

("v1", "/mnt/data/models/MOSS-Audio-Tokenizer/config.json", "moss_audio_tokenizer_v1_config"),

("v2", "/mnt/data/models/MOSS-Audio-Tokenizer-v2-cfg/config.json", "moss_audio_tokenizer_v2_config"),

("nano", "/mnt/data/models/MOSS-Audio-Tokenizer-Nano/config.json", "moss_audio_tokenizer_nano_config"),

]:

d = derive(cfg); h = hardcoded(src, fn)

print(f"=== {label}")

for side in ("encoder", "decoder"):

want, final, expect = d[side]

got = h[side]

ok = want == got

print(f" {side:8s} want {want}")

print(f" {'':8s} got {got} {'OK' if ok else '<-- MISMATCH'}")

if abs(final - expect) > 1:

print(f" {'':8s} frame-rate chain does not close: {final} vs {expect}")