chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,58 @@
|
||||
#!/usr/bin/env python3
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from examples.speech_recognition.data.collaters import Seq2SeqCollater
|
||||
|
||||
|
||||
class TestSeq2SeqCollator(unittest.TestCase):
|
||||
def test_collate(self):
|
||||
|
||||
eos_idx = 1
|
||||
pad_idx = 0
|
||||
collater = Seq2SeqCollater(
|
||||
feature_index=0, label_index=1, pad_index=pad_idx, eos_index=eos_idx
|
||||
)
|
||||
|
||||
# 2 frames in the first sample and 3 frames in the second one
|
||||
frames1 = np.array([[7, 8], [9, 10]])
|
||||
frames2 = np.array([[1, 2], [3, 4], [5, 6]])
|
||||
target1 = np.array([4, 2, 3, eos_idx])
|
||||
target2 = np.array([3, 2, eos_idx])
|
||||
sample1 = {"id": 0, "data": [frames1, target1]}
|
||||
sample2 = {"id": 1, "data": [frames2, target2]}
|
||||
batch = collater.collate([sample1, sample2])
|
||||
|
||||
# collate sort inputs by frame's length before creating the batch
|
||||
self.assertTensorEqual(batch["id"], torch.tensor([1, 0]))
|
||||
self.assertEqual(batch["ntokens"], 7)
|
||||
self.assertTensorEqual(
|
||||
batch["net_input"]["src_tokens"],
|
||||
torch.tensor(
|
||||
[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10], [pad_idx, pad_idx]]]
|
||||
),
|
||||
)
|
||||
self.assertTensorEqual(
|
||||
batch["net_input"]["prev_output_tokens"],
|
||||
torch.tensor([[eos_idx, 3, 2, pad_idx], [eos_idx, 4, 2, 3]]),
|
||||
)
|
||||
self.assertTensorEqual(batch["net_input"]["src_lengths"], torch.tensor([3, 2]))
|
||||
self.assertTensorEqual(
|
||||
batch["target"],
|
||||
torch.tensor([[3, 2, eos_idx, pad_idx], [4, 2, 3, eos_idx]]),
|
||||
)
|
||||
self.assertEqual(batch["nsentences"], 2)
|
||||
|
||||
def assertTensorEqual(self, t1, t2):
|
||||
self.assertEqual(t1.size(), t2.size(), "size mismatch")
|
||||
self.assertEqual(t1.ne(t2).long().sum(), 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user