1from __future__ import annotations

2

3import argparse

4import hashlib

5import inspect

6import json

7import random

8import shutil

9import sys

10import time

11from collections import defaultdict

12from pathlib import Path

13from typing import Any, Dict, List, Sequence, Tuple

14

15from ai_guitar_pedal.storage_redirect import (

16 E_ROOT,

17 apply_env_defaults,

18 assert_under_root,

19 ensure_storage_tree,

20 resolve_e_root,

21)

22

23

24def early_arg_value(argv: Sequence[str], name: str) -> str | None:

25 prefix = f"{name}="

26 for idx, item in enumerate(argv):

27 if item == name and idx + 1 < len(argv):

28 return argv[idx + 1]

29 if item.startswith(prefix):

30 return item[len(prefix) :]

31 return None

32

33

34def bootstrap_storage_env_from_argv(argv: Sequence[str]) -> None:

35 raw_e_root = early_arg_value(argv, "--e_root") or E_ROOT

36 try:

37 e_root = resolve_e_root(raw_e_root)

38 storage = ensure_storage_tree(e_root)

39 apply_env_defaults(storage, set_temp_tmp=True)

40 except Exception:

41 return

42

43

44bootstrap_storage_env_from_argv(sys.argv[1:])

45

46import numpy as np

47import torch

48from torch.utils.data import Dataset

49from transformers import (

50 AutoFeatureExtractor,

51 Trainer,

52 TrainingArguments,

53 WhisperConfig,

54 WhisperForConditionalGeneration,

55 set_seed,

56)

57

58from ai_guitar_pedal.memory_guard import MemoryWarningCallback

59from ai_guitar_pedal.tokenizer.dsp_control_tokenizer import DSPControlTokenizer

60from ai_guitar_pedal.phase2_training_utils import (

61 AUDIO_CHANNEL_MODE,

62 PauseSaveAndStopCallback,

63 STEREO_FEATURE_BINS,

64 TARGET_SR,

65 adapt_whisper_model_for_stereo,

66 feature_cache_base_dir,

67 feature_cache_path_for_audio,

68 load_or_compute_input_features,

69 preextract_feature_cache,

70 resolve_pause_paths,

71 resolve_resume_checkpoint,

72 sanitize_experiment_name,

73 str2bool,

74)

75

76

77def parse_args() -> argparse.Namespace:

78 parser = argparse.ArgumentParser(

79 description=(

80 "Train Whisper audio-to-DSP-control-token sequence model. "

81 "Targets are program model_token_ids, not preset numbers."

82 )

83 )

84 parser.add_argument("--e_root", default=E_ROOT)

85 parser.add_argument("--allow_non_e_root", type=str2bool, default=False)

86 parser.add_argument("--dry_run", type=str2bool, default=False)

87 parser.add_argument(

88 "--validate_only",

89 type=str2bool,

90 default=False,

91 help=(

92 "Load and validate manifest/program-token rows, write split manifests, "

93 "then exit before feature extraction/model training."

94 ),

95 )

96 parser.add_argument("--experiment_name", default="Phase2_ParamSeq_WhisperBase_V4_0001")

97 parser.add_argument("--model_name", default="openai/whisper-base")

98 parser.add_argument("--manifest", required=True)

99 parser.add_argument(

100 "--vocab",

101 default=str(

102 (

103 Path(__file__).resolve().parent

104 / "resources"

105 / "dsp_control_vocab.json"

106 ).resolve()

107 ),

108 )

109 parser.add_argument(

110 "--program_root",

111 default=str(

112 (

113 Path(__file__).resolve().parent

114 / ".."

115 / "Audio"

116 / "Presets"

117 / "programs"

118 ).resolve()

119 ),

120 help=(

121 "Local fallback directory for program JSON files. Used when "

122 "manifest target_program_json points at an unavailable absolute path."

123 ),

124 )

125 parser.add_argument("--audio_root", default="")

126 parser.add_argument("--feature_cache_root", default="")

127 parser.add_argument("--use_feature_cache", type=str2bool, default=True)

128 parser.add_argument("--preextract_features", type=str2bool, default=True)

129 parser.add_argument("--preextract_report_every", type=int, default=1000)

130 parser.add_argument("--output_dir", default="")

131 parser.add_argument("--log_dir", default="")

132 parser.add_argument("--export_dir", default="")

133 parser.add_argument("--max_seconds", type=float, default=6.0)

134 parser.add_argument("--max_target_length", type=int, default=768)

135 parser.add_argument("--strip_bos_from_labels", type=str2bool, default=True)

136 parser.add_argument("--eval_ratio", type=float, default=0.02)

137 parser.add_argument("--seed", type=int, default=42)

138 parser.add_argument("--per_device_train_batch_size", type=int, default=1)

139 parser.add_argument("--per_device_eval_batch_size", type=int, default=1)

140 parser.add_argument("--gradient_accumulation_steps", type=int, default=32)

141 parser.add_argument("--learning_rate", type=float, default=3e-5)

142 parser.add_argument("--num_train_epochs", type=float, default=3)

143 parser.add_argument("--max_steps", type=int, default=-1)

144 parser.add_argument("--fp16", type=str2bool, default=True)

145 parser.add_argument("--gradient_checkpointing", type=str2bool, default=True)

146 parser.add_argument("--gpu_mem_frac", type=float, default=0.90)

147 parser.add_argument("--num_workers", type=int, default=0)

148 parser.add_argument("--prefetch_factor", type=int, default=2)

149 parser.add_argument(

150 "--checkpoint_strategy",

151 choices=["epoch", "steps"],

152 default="steps",

153 )

154 parser.add_argument("--save_steps", type=int, default=250)

155 parser.add_argument("--eval_steps", type=int, default=250)

156 parser.add_argument("--save_total_limit", type=int, default=5)

157 parser.add_argument("--eval_accumulation_steps", type=int, default=1)

158 parser.add_argument("--auto_resume", type=str2bool, default=True)

159 parser.add_argument("--resume_from_checkpoint", default="")

160 parser.add_argument("--pause_flag_file", default="")

161 parser.add_argument("--pause_ckpt_dir", default="")

162 parser.add_argument("--mem_warn_enabled", type=str2bool, default=True)

163 parser.add_argument("--mem_warn_ram_threshold", type=float, default=0.90)

164 parser.add_argument("--mem_warn_vram_threshold", type=float, default=0.90)

165 parser.add_argument("--mem_warn_check_steps", type=int, default=25)

166 parser.add_argument("--mem_warn_repeat_steps", type=int, default=200)

167 return parser.parse_args()

168

169

170def resolve_paths(args: argparse.Namespace) -> Tuple[Path, Dict[str, Path], Dict[str, str], Path, Path, Path, Path, Path, Path, Path, Path]:

171 e_root = resolve_e_root(args.e_root)

172 storage = ensure_storage_tree(e_root)

173 env_effective = apply_env_defaults(storage, set_temp_tmp=True)

174

175 exp = sanitize_experiment_name(args.experiment_name)

176 manifest = Path(args.manifest).expanduser().resolve()

177 vocab = Path(args.vocab).expanduser().resolve()

178 program_root = Path(args.program_root).expanduser().resolve()

179 audio_root = (

180 Path(args.audio_root).expanduser().resolve()

181 if str(args.audio_root).strip()

182 else manifest.parent.resolve()

183 )

184 output_dir = (

185 Path(args.output_dir).expanduser().resolve()

186 if str(args.output_dir).strip()

187 else (storage["artifacts_checkpoints"] / exp).resolve()

188 )

189 log_dir = (

190 Path(args.log_dir).expanduser().resolve()

191 if str(args.log_dir).strip()

192 else (storage["artifacts_runs"] / exp).resolve()

193 )

194 export_dir = (

195 Path(args.export_dir).expanduser().resolve()

196 if str(args.export_dir).strip()

197 else (storage["artifacts_exports"] / exp).resolve()

198 )

199 feature_cache_root = (

200 Path(args.feature_cache_root).expanduser().resolve()

201 if str(args.feature_cache_root).strip()

202 else storage["cache_whisper_features"].resolve()

203 )

204

205 for path in [output_dir, log_dir, export_dir, feature_cache_root]:

206 path.mkdir(parents=True, exist_ok=True)

207

208 for name, path in [

209 ("output_dir", output_dir),

210 ("log_dir", log_dir),

211 ("export_dir", export_dir),

212 ("feature_cache_root", feature_cache_root),

213 ]:

214 assert_under_root(path, e_root, name, allow_non_root=bool(args.allow_non_e_root))

215

216 return (

217 e_root,

218 storage,

219 env_effective,

220 manifest,

221 vocab,

222 program_root,

223 audio_root,

224 output_dir,

225 log_dir,

226 export_dir,

227 feature_cache_root,

228 )

229

230

231def resolve_manifest_audio_path(manifest_dir: Path, raw_path: str) -> Path:

232 path = Path(str(raw_path).strip())

233 if path.is_absolute():

234 return path.expanduser().resolve()

235 return (manifest_dir / path).expanduser().resolve()

236

237

238def resolve_program_path(manifest_dir: Path, program_root: Path, raw_path: str) -> Path:

239 path = Path(str(raw_path).strip())

240 if path.is_absolute():

241 resolved = path.expanduser().resolve()

242 else:

243 resolved = (manifest_dir / path).expanduser().resolve()

244 if resolved.exists():

245 return resolved

246 fallback = (program_root / path.name).expanduser().resolve()

247 return fallback

248

249

250def token_hash(token_ids: Sequence[int]) -> str:

251 text = ",".join(str(int(x)) for x in token_ids)

252 return hashlib.sha256(text.encode("utf-8")).hexdigest()

253

254

255def read_program_token_ids(

256 program_path: Path,

257 tokenizer: DSPControlTokenizer,

258 strip_bos: bool,

259) -> List[int]:

260 payload = json.loads(program_path.read_text(encoding="utf-8"))

261 token_ids = payload.get("model_token_ids")

262 if not isinstance(token_ids, list) or not token_ids:

263 raise ValueError(f"program missing model_token_ids: {program_path}")

264 out = [int(x) for x in token_ids]

265 if strip_bos and out and out[0] == tokenizer.bos_token_id:

266 out = out[1:]

267 if not out or out[-1] != tokenizer.eos_token_id:

268 raise ValueError(f"program token sequence must end with EOS: {program_path}")

269 bad = [x for x in out if x < 0 or x >= tokenizer.vocab_size]

270 if bad:

271 raise ValueError(f"program token id out of vocabulary: {program_path}, examples={bad[:10]}")

272 return out

273

274

275def load_program_rows(

276 manifest_path: Path,

277 program_root: Path,

278 tokenizer: DSPControlTokenizer,

279 strip_bos: bool,

280 max_target_length: int,

281) -> List[Dict[str, Any]]:

282 manifest_path = manifest_path.expanduser().resolve()

283 manifest_dir = manifest_path.parent

284 program_cache: Dict[Path, Dict[str, Any]] = {}

285 rows: List[Dict[str, Any]] = []

286 invalid: List[str] = []

287 missing: List[str] = []

288

289 with manifest_path.open("r", encoding="utf-8") as handle:

290 for line_no, line in enumerate(handle, start=1):

291 text = line.strip()

292 if not text:

293 continue

294 try:

295 item = json.loads(text)

296 except json.JSONDecodeError as exc:

297 invalid.append(f"line {line_no}: invalid JSON ({exc})")

298 continue

299 if not isinstance(item, dict):

300 invalid.append(f"line {line_no}: expected object")

301 continue

302

303 raw_audio = item.get("audio_filepath")

304 raw_program = item.get("target_program_json")

305 if not isinstance(raw_audio, str) or not raw_audio.strip():

306 invalid.append(f"line {line_no}: missing audio_filepath")

307 continue

308 if not isinstance(raw_program, str) or not raw_program.strip():

309 invalid.append(f"line {line_no}: missing target_program_json")

310 continue

311

312 audio_path = resolve_manifest_audio_path(manifest_dir, raw_audio)

313 program_path = resolve_program_path(manifest_dir, program_root, raw_program)

314 if not audio_path.exists():

315 missing.append(f"line {line_no}: audio not found: {audio_path}")

316 continue

317 if not program_path.exists():

318 missing.append(f"line {line_no}: program not found: {program_path}")

319 continue

320

321 if program_path not in program_cache:

322 token_ids = read_program_token_ids(

323 program_path=program_path,

324 tokenizer=tokenizer,

325 strip_bos=strip_bos,

326 )

327 if len(token_ids) > max_target_length:

328 invalid.append(

329 f"line {line_no}: target length {len(token_ids)} > max_target_length "

330 f"{max_target_length}: {program_path}"

331 )

332 continue

333 program_cache[program_path] = {

334 "labels": token_ids,

335 "program_hash": token_hash(token_ids),

336 "target_length": len(token_ids),

337 }

338

339 cached = program_cache[program_path]

340 rows.append(

341 {

342 "audio": str(audio_path),

343 "labels": cached["labels"],

344 "program_hash": cached["program_hash"],

345 "target_length": cached["target_length"],

346 }

347 )

348

349 if invalid or missing:

350 lines = [f"Invalid program training manifest: {manifest_path}"]

351 if invalid:

352 lines.append(f"- Field/token errors: {len(invalid)}")

353 lines.extend(f" {x}" for x in invalid[:100])

354 if len(invalid) > 100:

355 lines.append(f" ... and {len(invalid) - 100} more")

356 if missing:

357 lines.append(f"- Missing files: {len(missing)}")

358 lines.extend(f" {x}" for x in missing[:100])

359 if len(missing) > 100:

360 lines.append(f" ... and {len(missing) - 100} more")

361 raise ValueError("\n".join(lines))

362 if not rows:

363 raise ValueError(f"No usable rows found in manifest: {manifest_path}")

364 return rows

365

366

367def split_by_program_hash(

368 rows: Sequence[Dict[str, Any]],

369 eval_ratio: float,

370 seed: int,

371) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:

372 if not (0.0 <= eval_ratio < 1.0):

373 raise ValueError(f"eval_ratio must be in [0,1), got {eval_ratio}")

374 grouped: Dict[str, List[Dict[str, Any]]] = defaultdict(list)

375 for row in rows:

376 grouped[str(row["program_hash"])].append(row)

377

378 keys = sorted(grouped)

379 rng = random.Random(seed)

380 rng.shuffle(keys)

381 if eval_ratio <= 0.0 or len(keys) <= 1:

382 eval_keys = set()

383 else:

384 eval_count = int(round(len(keys) * eval_ratio))

385 eval_count = max(1, eval_count)

386 eval_count = min(eval_count, len(keys) - 1)

387 eval_keys = set(keys[:eval_count])

388

389 train_rows: List[Dict[str, Any]] = []

390 eval_rows: List[Dict[str, Any]] = []

391 for key in keys:

392 if key in eval_keys:

393 eval_rows.extend(grouped[key])

394 else:

395 train_rows.extend(grouped[key])

396 rng.shuffle(train_rows)

397 rng.shuffle(eval_rows)

398 if not train_rows:

399 raise ValueError("No train rows after program-hash split.")

400 return train_rows, eval_rows

401

402

403def write_split_manifest(path: Path, rows: Sequence[Dict[str, Any]]) -> None:

404 path.parent.mkdir(parents=True, exist_ok=True)

405 with path.open("w", encoding="utf-8") as handle:

406 for row in rows:

407 handle.write(

408 json.dumps(

409 {

410 "audio": row["audio"],

411 "program_hash": row["program_hash"],

412 "target_length": int(row["target_length"]),

413 },

414 ensure_ascii=False,

415 )

416 + "\n"

417 )

418

419

420class ProgramSeq2SeqDataset(Dataset):

421 def __init__(

422 self,

423 rows: Sequence[Dict[str, Any]],

424 feature_extractor: AutoFeatureExtractor,

425 max_seconds: float,

426 audio_root: Path,

427 use_feature_cache: bool,

428 cache_base: Path,

429 ) -> None:

430 self.rows = list(rows)

431 self.feature_extractor = feature_extractor

432 self.max_len = int(max_seconds * TARGET_SR)

433 self.audio_root = audio_root

434 self.use_feature_cache = bool(use_feature_cache)

435 self.cache_base = cache_base

436 if self.max_len <= 0:

437 raise ValueError(f"max_seconds must produce positive max_len, got {max_seconds}")

438

439 def __len__(self) -> int:

440 return len(self.rows)

441

442 def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:

443 row = self.rows[idx]

444 wav_path = Path(str(row["audio"])).expanduser().resolve()

445 cache_path = feature_cache_path_for_audio(

446 wav_path=wav_path,

447 audio_root=self.audio_root,

448 cache_base=self.cache_base,

449 )

450 input_features = load_or_compute_input_features(

451 wav_path=wav_path,

452 feature_extractor=self.feature_extractor,

453 max_len=self.max_len,

454 use_feature_cache=self.use_feature_cache,

455 cache_path=cache_path,

456 save_on_miss=self.use_feature_cache,

457 )

458 return {

459 "input_features": input_features,

460 "labels": torch.tensor(row["labels"], dtype=torch.long),

461 }

462

463

464class Seq2SeqCollator:

465 def __init__(self, pad_token_id: int, label_pad_token_id: int = -100) -> None:

466 self.pad_token_id = int(pad_token_id)

467 self.label_pad_token_id = int(label_pad_token_id)

468

469 def __call__(self, items: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:

470 input_features = torch.stack([item["input_features"] for item in items], dim=0)

471 labels = [item["labels"] for item in items]

472 max_len = max(int(label.shape[0]) for label in labels)

473 padded = torch.full(

474 (len(labels), max_len),

475 self.label_pad_token_id,

476 dtype=torch.long,

477 )

478 for idx, label in enumerate(labels):

479 padded[idx, : label.shape[0]] = label

480 return {"input_features": input_features, "labels": padded}

481

482

483def preprocess_logits_for_metrics(logits, labels):

484 if isinstance(logits, tuple):

485 logits = logits[0]

486 return torch.argmax(logits, dim=-1)

487

488

489def compute_token_metrics(eval_pred) -> Dict[str, float]:

490 preds = np.asarray(eval_pred.predictions)

491 labels = np.asarray(eval_pred.label_ids)

492 if preds.shape != labels.shape:

493 preds = preds.reshape(labels.shape)

494 mask = labels != -100

495 if not np.any(mask):

496 return {"token_acc": 0.0, "seq_exact": 0.0}

497 token_acc = float(np.mean(preds[mask] == labels[mask]))

498 seq_ok = []

499 for pred_row, label_row, mask_row in zip(preds, labels, mask):

500 seq_ok.append(bool(np.all(pred_row[mask_row] == label_row[mask_row])))

501 return {

502 "token_acc": token_acc,

503 "seq_exact": float(np.mean(seq_ok)) if seq_ok else 0.0,

504 }

505

506

507def build_model(args: argparse.Namespace, tokenizer: DSPControlTokenizer) -> WhisperForConditionalGeneration:

508 config = WhisperConfig.from_pretrained(args.model_name)

509 requested_target_positions = int(args.max_target_length)

510 if requested_target_positions <= 0:

511 raise ValueError(f"max_target_length must be positive, got {requested_target_positions}")

512 config.max_target_positions = max(

513 requested_target_positions,

514 int(getattr(config, "max_target_positions", 0) or 0),

515 )

516 config.vocab_size = int(tokenizer.vocab_size)

517 config.pad_token_id = int(tokenizer.pad_token_id)

518 config.bos_token_id = int(tokenizer.bos_token_id)

519 config.eos_token_id = int(tokenizer.eos_token_id)

520 config.decoder_start_token_id = int(tokenizer.bos_token_id)

521 config.forced_decoder_ids = None

522 config.suppress_tokens = []

523 config.begin_suppress_tokens = []

524

525 model = WhisperForConditionalGeneration.from_pretrained(

526 args.model_name,

527 config=config,

528 ignore_mismatched_sizes=True,

529 )

530 adapt_whisper_model_for_stereo(model)

531 model.config.use_cache = False

532 if hasattr(model, "generation_config"):

533 model.generation_config.pad_token_id = int(tokenizer.pad_token_id)

534 model.generation_config.bos_token_id = int(tokenizer.bos_token_id)

535 model.generation_config.eos_token_id = int(tokenizer.eos_token_id)

536 model.generation_config.decoder_start_token_id = int(tokenizer.bos_token_id)

537 model.generation_config.forced_decoder_ids = None

538 model.generation_config.suppress_tokens = []

539 model.generation_config.begin_suppress_tokens = []

540 model.generation_config.max_length = int(args.max_target_length)

541 if args.gradient_checkpointing:

542 model.gradient_checkpointing_enable()

543 return model

544

545

546def build_training_args(args: argparse.Namespace, has_eval: bool) -> TrainingArguments:

547 strategy = str(args.checkpoint_strategy)

548 if strategy == "steps":

549 if int(args.save_steps) <= 0:

550 raise ValueError("save_steps must be > 0")

551 eval_steps = int(args.eval_steps) if int(args.eval_steps) > 0 else int(args.save_steps)

552 if int(args.save_steps) % int(eval_steps) != 0:

553 raise ValueError("save_steps must be a round multiple of eval_steps")

554 else:

555 eval_steps = 0

556

557 kwargs: Dict[str, Any] = {

558 "output_dir": str(Path(args.output_dir).expanduser().resolve()),

559 "logging_dir": str(Path(args.log_dir).expanduser().resolve()),

560 "per_device_train_batch_size": int(args.per_device_train_batch_size),

561 "per_device_eval_batch_size": int(args.per_device_eval_batch_size),

562 "gradient_accumulation_steps": int(args.gradient_accumulation_steps),

563 "learning_rate": float(args.learning_rate),

564 "num_train_epochs": float(args.num_train_epochs),

565 "max_steps": int(args.max_steps),

566 "fp16": bool(args.fp16 and torch.cuda.is_available()),

567 "save_strategy": strategy,

568 "logging_strategy": "steps",

569 "logging_steps": 10,

570 "save_total_limit": max(1, int(args.save_total_limit)),

571 "report_to": [],

572 "run_name": sanitize_experiment_name(args.experiment_name),

573 "remove_unused_columns": False,

574 "seed": int(args.seed),

575 "gradient_checkpointing": bool(args.gradient_checkpointing),

576 "dataloader_num_workers": int(args.num_workers),

577 "dataloader_pin_memory": True,

578 "eval_accumulation_steps": max(1, int(args.eval_accumulation_steps)),

579 }

580 if int(args.max_steps) <= 0:

581 kwargs.pop("max_steps")

582

583 params = inspect.signature(TrainingArguments.__init__).parameters

584 if "eval_strategy" in params:

585 kwargs["eval_strategy"] = strategy if has_eval else "no"

586 else:

587 kwargs["evaluation_strategy"] = strategy if has_eval else "no"

588 if strategy == "steps":

589 kwargs["save_steps"] = int(args.save_steps)

590 if has_eval:

591 kwargs["eval_steps"] = int(eval_steps)

592 if has_eval:

593 kwargs["load_best_model_at_end"] = True

594 kwargs["metric_for_best_model"] = "eval_token_acc"

595 kwargs["greater_is_better"] = True

596 if "dataloader_persistent_workers" in params:

597 kwargs["dataloader_persistent_workers"] = bool(args.num_workers > 0)

598 if "dataloader_prefetch_factor" in params and args.num_workers > 0:

599 kwargs["dataloader_prefetch_factor"] = max(1, int(args.prefetch_factor))

600 return TrainingArguments(**kwargs)

601

602

603def main() -> None:

604 args = parse_args()

605 (

606 e_root,

607 storage,

608 env_effective,

609 manifest_path,

610 vocab_path,

611 program_root,

612 audio_root,

613 output_dir,

614 log_dir,

615 export_dir,

616 feature_cache_root,

617 ) = resolve_paths(args)

618

619 args.manifest = str(manifest_path)

620 args.vocab = str(vocab_path)

621 args.program_root = str(program_root)

622 args.audio_root = str(audio_root)

623 args.output_dir = str(output_dir)

624 args.log_dir = str(log_dir)

625 args.export_dir = str(export_dir)

626 args.feature_cache_root = str(feature_cache_root)

627

628 print("[ParamSeq] Effective paths")

629 print(f" e_root={e_root}")

630 print(f" experiment_name={sanitize_experiment_name(args.experiment_name)}")

631 print(f" manifest={manifest_path}")

632 print(f" vocab={vocab_path}")

633 print(f" program_root={program_root}")

634 print(f" audio_root={audio_root}")

635 print(f" output_dir={output_dir}")

636 print(f" log_dir={log_dir}")

637 print(f" export_dir={export_dir}")

638 print(f" feature_cache_root={feature_cache_root}")

639

640 if args.dry_run:

641 print("[DryRun] Exiting before dataset scan and training.")

642 return

643

644 if not manifest_path.exists():

645 raise FileNotFoundError(f"manifest not found: {manifest_path}")

646 if not vocab_path.exists():

647 raise FileNotFoundError(f"vocab not found: {vocab_path}")

648 if not program_root.exists():

649 raise FileNotFoundError(f"program_root not found: {program_root}")

650 if not audio_root.exists():

651 raise FileNotFoundError(f"audio_root not found: {audio_root}")

652

653 set_seed(int(args.seed))

654 tokenizer = DSPControlTokenizer.from_vocab_file(vocab_path)

655

656 pause_flag_file, pause_ckpt_dir = resolve_pause_paths(args, output_dir)

657 assert_under_root(

658 pause_ckpt_dir,

659 e_root,

660 "pause_ckpt_dir",

661 allow_non_root=bool(args.allow_non_e_root),

662 )

663

664 if args.gpu_mem_frac and torch.cuda.is_available():

665 if not (0.0 < float(args.gpu_mem_frac) <= 1.0):

666 raise ValueError(f"gpu_mem_frac must be in (0,1], got {args.gpu_mem_frac}")

667 if hasattr(torch.cuda, "memory") and hasattr(

668 torch.cuda.memory, "set_per_process_memory_fraction"

669 ):

670 torch.cuda.memory.set_per_process_memory_fraction(float(args.gpu_mem_frac), device=0)

671 else:

672 torch.cuda.set_per_process_memory_fraction(float(args.gpu_mem_frac), device=0)

673 total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3

674 print(

675 f"[GPU] total={total_mem_gb:.2f}GB, "

676 f"limit={float(args.gpu_mem_frac):.2f}, "

677 f"allowed~{total_mem_gb * float(args.gpu_mem_frac):.2f}GB"

678 )

679

680 t0 = time.time()

681 rows = load_program_rows(

682 manifest_path=manifest_path,

683 program_root=program_root,

684 tokenizer=tokenizer,

685 strip_bos=bool(args.strip_bos_from_labels),

686 max_target_length=int(args.max_target_length),

687 )

688 elapsed = time.time() - t0

689 lengths = [int(row["target_length"]) for row in rows]

690 program_hash_count = len({str(row["program_hash"]) for row in rows})

691 print(

692 f"[ParamSeq] loaded rows={len(rows)}, unique_program_hashes={program_hash_count}, "

693 f"target_len_min={min(lengths)}, target_len_max={max(lengths)}, "

694 f"target_len_mean={float(np.mean(lengths)):.2f}, elapsed={elapsed:.1f}s"

695 )

696

697 train_rows, eval_rows = split_by_program_hash(rows, float(args.eval_ratio), int(args.seed))

698 print(

699 f"[ParamSeq] split by program_hash: train_rows={len(train_rows)}, "

700 f"eval_rows={len(eval_rows)}, train_programs={len({r['program_hash'] for r in train_rows})}, "

701 f"eval_programs={len({r['program_hash'] for r in eval_rows})}"

702 )

703 write_split_manifest(output_dir / "train_split_desensitized.jsonl", train_rows)

704 write_split_manifest(output_dir / "eval_split_desensitized.jsonl", eval_rows)

705 if args.validate_only:

706 print("[ValidateOnly] Manifest/program-token validation passed.")

707 print(f"[ValidateOnly] Split manifests written under: {output_dir}")

708 return

709

710 feature_extractor = AutoFeatureExtractor.from_pretrained(args.model_name)

711 print(f"Audio channel mode: {AUDIO_CHANNEL_MODE}")

712 print(f"Audio feature bins: {STEREO_FEATURE_BINS}")

713 cache_base = feature_cache_base_dir(

714 feature_cache_root=feature_cache_root,

715 model_name=args.model_name,

716 max_seconds=float(args.max_seconds),

717 )

718 if args.use_feature_cache and args.preextract_features:

719 preextract_feature_cache(

720 rows=rows,

721 audio_root=audio_root,

722 feature_extractor=feature_extractor,

723 max_len=int(float(args.max_seconds) * TARGET_SR),

724 cache_base=cache_base,

725 report_every=int(args.preextract_report_every),

726 )

727

728 train_dataset = ProgramSeq2SeqDataset(

729 rows=train_rows,

730 feature_extractor=feature_extractor,

731 max_seconds=float(args.max_seconds),

732 audio_root=audio_root,

733 use_feature_cache=bool(args.use_feature_cache),

734 cache_base=cache_base,

735 )

736 eval_dataset = (

737 ProgramSeq2SeqDataset(

738 rows=eval_rows,

739 feature_extractor=feature_extractor,

740 max_seconds=float(args.max_seconds),

741 audio_root=audio_root,

742 use_feature_cache=bool(args.use_feature_cache),

743 cache_base=cache_base,

744 )

745 if eval_rows

746 else None

747 )

748

749 model = build_model(args, tokenizer)

750 trainer = Trainer(

751 model=model,

752 args=build_training_args(args, has_eval=eval_dataset is not None),

753 train_dataset=train_dataset,

754 eval_dataset=eval_dataset,

755 data_collator=Seq2SeqCollator(pad_token_id=tokenizer.pad_token_id),

756 compute_metrics=compute_token_metrics if eval_dataset is not None else None,

757 preprocess_logits_for_metrics=preprocess_logits_for_metrics

758 if eval_dataset is not None

759 else None,

760 )

761 if args.mem_warn_enabled:

762 trainer.add_callback(

763 MemoryWarningCallback(

764 ram_threshold=float(args.mem_warn_ram_threshold),

765 vram_threshold=float(args.mem_warn_vram_threshold),

766 check_every_steps=int(args.mem_warn_check_steps),

767 repeat_warning_steps=int(args.mem_warn_repeat_steps),

768 )

769 )

770

771 pause_callback = PauseSaveAndStopCallback(

772 pause_flag_file=pause_flag_file,

773 pause_ckpt_dir=pause_ckpt_dir,

774 remove_pause_flag=True,

775 )

776 pause_callback.set_trainer(trainer)

777 trainer.add_callback(pause_callback)

778

779 resume_ckpt, resume_source = resolve_resume_checkpoint(args, pause_ckpt_dir)

780 if resume_ckpt is not None:

781 print(f"[Resume] source={resume_source}, resume_from_checkpoint={resume_ckpt}")

782 trainer.train(resume_from_checkpoint=str(resume_ckpt))

783 else:

784 print(f"[Resume] source={resume_source}, starting from scratch")

785 trainer.train()

786

787 if eval_dataset is not None:

788 metrics = trainer.evaluate()

789 print(

790 f"[Eval] eval_token_acc={metrics.get('eval_token_acc', 0.0):.6f}, "

791 f"eval_seq_exact={metrics.get('eval_seq_exact', 0.0):.6f}"

792 )

793 else:

794 metrics = {}

795

796 export_dir.mkdir(parents=True, exist_ok=True)

797 trainer.save_model(str(export_dir))

798 feature_extractor.save_pretrained(str(export_dir))

799 shutil.copy2(vocab_path, export_dir / "dsp_control_vocab.json")

800

801 summary = {

802 "objective": "audio_to_dsp_control_program_tokens",

803 "desensitized_from_preset_number": True,

804 "uses_preset_number_as_label": False,

805 "target_source_field": "target_program_json.model_token_ids",

806 "manifest": str(manifest_path),

807 "vocab": str(vocab_path),

808 "program_root": str(program_root),

809 "vocab_size": int(tokenizer.vocab_size),

810 "audio_root": str(audio_root),

811 "total_rows": int(len(rows)),

812 "train_rows": int(len(train_rows)),

813 "eval_rows": int(len(eval_rows)),

814 "unique_program_hashes": int(program_hash_count),

815 "train_program_hashes": int(len({r["program_hash"] for r in train_rows})),

816 "eval_program_hashes": int(len({r["program_hash"] for r in eval_rows})),

817 "target_length": {

818 "min": int(min(lengths)),

819 "max": int(max(lengths)),

820 "mean": float(np.mean(lengths)),

821 },

822 "metrics": metrics,

823 "training_config": {

824 "model_name": str(args.model_name),

825 "audio_channel_mode": AUDIO_CHANNEL_MODE,

826 "input_feature_bins": int(STEREO_FEATURE_BINS),

827 "max_seconds": float(args.max_seconds),

828 "max_target_length": int(args.max_target_length),

829 "strip_bos_from_labels": bool(args.strip_bos_from_labels),

830 "eval_ratio": float(args.eval_ratio),

831 "num_train_epochs": float(args.num_train_epochs),

832 "max_steps": int(args.max_steps),

833 "checkpoint_strategy": str(args.checkpoint_strategy),

834 "save_steps": int(args.save_steps),

835 "eval_steps": int(args.eval_steps),

836 "save_total_limit": int(args.save_total_limit),

837 "per_device_train_batch_size": int(args.per_device_train_batch_size),

838 "per_device_eval_batch_size": int(args.per_device_eval_batch_size),

839 "gradient_accumulation_steps": int(args.gradient_accumulation_steps),

840 "learning_rate": float(args.learning_rate),

841 "gradient_checkpointing": bool(args.gradient_checkpointing),

842 "fp16": bool(args.fp16 and torch.cuda.is_available()),

843 "gpu_mem_frac": float(args.gpu_mem_frac),

844 "auto_resume": bool(args.auto_resume),

845 "resume_from_checkpoint": str(resume_ckpt) if resume_ckpt is not None else "",

846 "resume_source": str(resume_source),

847 },

848 "feature_cache_base": str(cache_base),

849 "e_root": str(e_root),

850 "output_dir": str(output_dir),

851 "log_dir": str(log_dir),

852 "export_dir": str(export_dir),

853 "pause_flag_file": str(pause_flag_file),

854 "pause_ckpt_dir": str(pause_ckpt_dir),

855 "env": env_effective,

856 }

857 (export_dir / "run_summary.json").write_text(

858 json.dumps(summary, indent=2, ensure_ascii=False),

859 encoding="utf-8",

860 )

861 print(f"Parameter sequence model exported to: {export_dir}")

862

863

864if __name__ == "__main__":

865 main()