Skip to content

Instantly share code, notes, and snippets.

@zilunpeng
zilunpeng / model_data_path.py
Last active March 23, 2021 21:00
Specify paths to model and dataset. Code below is part of the wav2vec 2.0 inference notebook (https://git.io/JYeKX).
model_path = "/home/models/wav2vec_big_960h.pt"
data_path = "/home/datasets/"
@zilunpeng
zilunpeng / create_target_dict.py
Created March 23, 2021 21:01
Create mapping between tokens and indices. Code below is part of the wav2vec 2.0 inference notebook (https://git.io/JYeKX).
target_dict = fairseq_mod.data.Dictionary.load('ltr_dict.txt')
@zilunpeng
zilunpeng / create_wav2vec2.py
Created March 23, 2021 21:03
Create and load weights into wav2vec 2.0. Code below is part of the wav2vec 2.0 inference notebook (https://git.io/JYeKX).
w2v = torch.load(model_path)
model = Wav2VecCtc.build_model(w2v["args"], target_dict)
model.load_state_dict(w2v["model"], strict=True)
@zilunpeng
zilunpeng / create_decoder.py
Created March 23, 2021 21:04
Create a Viterbi decoder. Code below is part of the wav2vec 2.0 inference notebook (https://git.io/JYeKX).
decoder = W2lViterbiDecoder(target_dict)
@zilunpeng
zilunpeng / create_dev_clean_data_loader.py
Created March 23, 2021 21:05
Create data loader. Code below is part of the wav2vec 2.0 inference notebook (https://git.io/JYeKX).
dev_clean_librispeech_data = torchaudio.datasets.LIBRISPEECH(data_path, url='dev-clean', download=False)
data_loader = torch.utils.data.DataLoader(dev_clean_librispeech_data, batch_size=1, shuffle=False)
@zilunpeng
zilunpeng / get_wav2vec2_output.py
Created March 23, 2021 21:07
Get the output from wav2vec 2.0. Code below is part of the wav2vec 2.0 inference notebook (https://git.io/JYeKX).
encoder_out = model(**encoder_input)
emissions = model.get_normalized_probs(encoder_out, log_probs=True)
emissions = emissions.transpose(0, 1).float().cpu().contiguous()
@zilunpeng
zilunpeng / get_wav2vec2_decoder_output.py
Created March 23, 2021 21:08
Get output from decoder. Code below is part of the wav2vec 2.0 inference notebook (https://git.io/JYeKX).
decoder_out = decoder.decode(emissions)
@zilunpeng
zilunpeng / set_teacher_wav2vec2.py
Created March 23, 2021 21:15
Set the teacher model to evaluation mode. Code below is part of the knowledge distillation toolkit (https://git.io/JYePf).
self.teacher_model.eval()
@zilunpeng
zilunpeng / get_teacher_wav2vec2_prob.py
Created March 23, 2021 21:16
Get teacher model's probability distribution. Code below is part of the knowledge distillation toolkit (https://git.io/JYePf).
with torch.no_grad():
teacher_net_output = self.teacher_model(*batch)
teacher_prob = teacher_net_output["prob"]
@zilunpeng
zilunpeng / get_student_wav2vec2_log_prob.py
Created March 23, 2021 21:18
Get log probability of student model. Code below is part of the knowledge distillation toolkit (https://git.io/JYePf).
student_net_output = self.student_model(*batch)
student_log_prob = student_net_output["log_prob"]