()
| 50 | # Main evaluation routine |
| 51 | # --------------------------------------------------------------------- |
| 52 | def main() -> None: |
| 53 | args = parse_args() |
| 54 | |
| 55 | # -------- SentencePiece ---------- |
| 56 | sp = spm.SentencePieceProcessor() |
| 57 | sp.load(args.sp_model) |
| 58 | |
| 59 | # expose the tokenizer inside the imported `train` module so that |
| 60 | # train.int_to_text() works correctly during decoding |
| 61 | train.sp = sp |
| 62 | |
| 63 | vocab_size = sp.get_piece_size() + 1 # +1 for CTC blank |
| 64 | |
| 65 | # -------- Model ---------- |
| 66 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 67 | model = create_model(vocab_size).to(device) |
| 68 | |
| 69 | ckpt_path = Path(args.checkpoint) |
| 70 | ckpt = torch.load(ckpt_path, map_location=device) |
| 71 | state = ckpt.get("model", ckpt) # handles raw state_dict vs wrapper |
| 72 | |
| 73 | # tolerate DataParallel prefix differences |
| 74 | missing, unexpected = model.load_state_dict(state, strict=False) |
| 75 | if missing: |
| 76 | print(f"[warn] missing keys in checkpoint: {missing}") |
| 77 | if unexpected: |
| 78 | print(f"[warn] unexpected keys in checkpoint: {unexpected}") |
| 79 | |
| 80 | print(f"=> loaded weights from {ckpt_path}") |
| 81 | |
| 82 | if torch.cuda.device_count() > 1: |
| 83 | print(f"Using {torch.cuda.device_count()} GPUs via DataParallel …") |
| 84 | model = torch.nn.DataParallel(model) |
| 85 | |
| 86 | loss_fn = nn.CTCLoss(blank=0, zero_infinity=True) |
| 87 | |
| 88 | # -------- Datasets & evaluation ---------- |
| 89 | subsets = [s.strip() for s in args.set.split(",") if s.strip()] |
| 90 | preproc = AudioPreprocessor(training=False) |
| 91 | for subset in subsets: |
| 92 | ds = LIBRISPEECH(args.root, url=subset, download=True) |
| 93 | loader = DataLoader( |
| 94 | ds, |
| 95 | batch_size=args.batch_size, |
| 96 | shuffle=False, |
| 97 | collate_fn=collate_eval_factory(preproc), |
| 98 | num_workers=args.num_workers, |
| 99 | ) |
| 100 | |
| 101 | wer, val_loss = evaluate(model, loader, device, loss_fn) |
| 102 | print(f"\n── Results on {subset} ──") |
| 103 | print(f" • WER: {wer:6.2f} %") |
| 104 | print(f" • CTC loss: {val_loss:8.4f}") |
| 105 | |
| 106 | if __name__ == "__main__": |
| 107 | main() |
no test coverage detected