def test_sequence_reader(sequences, use_vocab, add_bos, add_eos):
with TemporaryDirectory() as work_dir:
path = os.path.join(work_dir, 'input')
with open(path, 'w') as f:
for sequence in sequences:
print(sequence, file=f)
vocabulary = vocab.build_vocab(sequences) if use_vocab else None
reader = data_io.SequenceReader(path, vocabulary=vocabulary, add_bos=add_bos, add_eos=add_eos)
read_sequences = [s for s in reader]
assert len(read_sequences) == len(sequences)
if vocabulary is None:
with pytest.raises(SockeyeError) as e:
_ = data_io.SequenceReader(path, vocabulary=vocabulary, add_bos=True)
assert str(e.value) == "Adding a BOS or EOS symbol requires a vocabulary"
expected_sequences = [data_io.strids2ids(get_tokens(s)) if s else None for s in sequences]
assert read_sequences == expected_sequences
else:
expected_sequences = [data_io.tokens2ids(get_tokens(s), vocabulary) if s else None for s in sequences]
if add_bos:
expected_sequences = [[vocabulary[C.BOS_SYMBOL]] + s if s else None for s in expected_sequences]
if add_eos:
expected_sequences = [s + [vocabulary[C.EOS_SYMBOL]] if s else None for s in expected_sequences]
assert read_sequences == expected_sequences
def test_tokens2ids(tokens, vocab, expected_ids):
ids = data_io.tokens2ids(tokens, vocab)
assert ids == expected_ids
def test_tokens2ids(tokens, vocab, expected_ids):
ids = data_io.tokens2ids(tokens, vocab)
assert ids == expected_ids
def test_tokens2ids(tokens, vocab, expected_ids):
pytest.importorskip('mxnet')
from sockeye import data_io
ids = data_io.tokens2ids(tokens, vocab)
assert ids == expected_ids