Commit cd4d6607 authored by hwangjeff's avatar hwangjeff Committed by Facebook GitHub Bot
Browse files

Fix bug in Conformer RNN-T recipe (#2611)

Summary:
https://github.com/pytorch/audio/issues/2535 modified the Conformer RNN-T Lightning module to accept a SentencePiece model instance rather than a file path. This PR makes changes to account for this in the train script.

Pull Request resolved: https://github.com/pytorch/audio/pull/2611

Reviewed By: carolineechen

Differential Revision: D38578892

Pulled By: hwangjeff

fbshipit-source-id: ec3b9823ad30ffb730baa13d10d8b79020866aac
parent fb4eb981
import pathlib
from argparse import ArgumentParser
import sentencepiece as spm
from lightning import ConformerRNNTModule
from pytorch_lightning import seed_everything, Trainer
from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
......@@ -44,7 +46,8 @@ def run_train(args):
reload_dataloaders_every_n_epochs=1,
)
model = ConformerRNNTModule(str(args.sp_model_path))
sp_model = spm.SentencePieceProcessor(model_file=str(args.sp_model_path))
model = ConformerRNNTModule(sp_model)
data_module = get_data_module(str(args.librispeech_path), str(args.global_stats_path), str(args.sp_model_path))
trainer.fit(model, data_module, ckpt_path=args.checkpoint_path)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment