← 返回训练指南

Main/ai_guitar_pedal/tokenizer/dsp_control_tokenizer.py

当前项目源码快照 · SHA-256 5fff29a5632d54f741ccae7450116d98ac0c3c6b20322552fc947fd6409f28f1

1from __future__ import annotations
2
3import json
4from pathlib import Path
5from typing import Iterable
6
7
8class DSPControlTokenizer:
9 """Strict finite-vocabulary tokenizer for DSP control programs."""
10
11 def __init__(self, token_to_id: dict[str, int]):
12 if not token_to_id:
13 raise ValueError("DSP control vocabulary is empty")
14 ids = list(token_to_id.values())
15 if len(ids) != len(set(ids)):
16 raise ValueError("DSP control vocabulary contains duplicate token IDs")
17 if sorted(ids) != list(range(len(ids))):
18 raise ValueError("DSP control token IDs must be contiguous from zero")
19
20 self.token_to_id = dict(token_to_id)
21 self.id_to_token = {token_id: token for token, token_id in token_to_id.items()}
22 self.vocab_size = len(token_to_id)
23 self.pad_token_id = self._required_id("<PAD>")
24 self.bos_token_id = self._required_id("<BOS>")
25 self.eos_token_id = self._required_id("<EOS>")
26
27 def _required_id(self, token: str) -> int:
28 if token not in self.token_to_id:
29 raise ValueError(f"Required DSP control token is missing: {token}")
30 return self.token_to_id[token]
31
32 def encode_tokens(self, tokens: Iterable[str]) -> list[int]:
33 ids = []
34 for token in tokens:
35 if token not in self.token_to_id:
36 raise ValueError(f"Unknown DSP control token: {token}")
37 ids.append(self.token_to_id[token])
38 return ids
39
40 def encode_text(self, text: str) -> list[int]:
41 return self.encode_tokens(text.split())
42
43 def decode_ids(self, token_ids: Iterable[int]) -> list[str]:
44 tokens = []
45 for token_id in token_ids:
46 token_id = int(token_id)
47 if token_id not in self.id_to_token:
48 raise ValueError(f"Unknown DSP control token ID: {token_id}")
49 tokens.append(self.id_to_token[token_id])
50 return tokens
51
52 def decode_text(self, token_ids: Iterable[int]) -> str:
53 return " ".join(self.decode_ids(token_ids))
54
55 @classmethod
56 def from_vocab_file(cls, path: str | Path) -> "DSPControlTokenizer":
57 payload = json.loads(Path(path).read_text(encoding="utf-8"))
58 return cls({token: int(token_id) for token, token_id in payload["token_to_id"].items()})
59