Spaces:
Paused
Paused
| #!/usr/bin/python | |
| # -*- coding: utf-8 -*- | |
| # Copyright 2018 Google Inc. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License.! | |
| import sys | |
| sys.path.insert(0, 'src') | |
| from collections import defaultdict | |
| import io | |
| import os | |
| import pickle | |
| import unittest | |
| import sentencepiece as spm | |
| print('VERSION={}'.format(spm.__version__)) | |
| data_dir = 'test' | |
| if sys.platform == 'win32': | |
| data_dir = os.path.join('..', 'data') | |
| class TestSentencepieceProcessor(unittest.TestCase): | |
| """Test case for SentencePieceProcessor""" | |
| def setUp(self): | |
| self.sp_ = spm.SentencePieceProcessor() | |
| self.jasp_ = spm.SentencePieceProcessor() | |
| self.assertTrue(self.sp_.Load(os.path.join('test', 'test_model.model'))) | |
| self.assertTrue( | |
| self.jasp_.Load(os.path.join('test', 'test_ja_model.model')) | |
| ) | |
| with open(os.path.join('test', 'test_model.model'), 'rb') as f: | |
| self.assertTrue(self.sp_.LoadFromSerializedProto(f.read())) | |
| with open(os.path.join('test', 'test_ja_model.model'), 'rb') as f: | |
| self.assertTrue(self.jasp_.LoadFromSerializedProto(f.read())) | |
| def test_load(self): | |
| self.assertEqual(1000, self.sp_.GetPieceSize()) | |
| self.assertEqual(0, self.sp_.PieceToId('<unk>')) | |
| self.assertEqual(1, self.sp_.PieceToId('<s>')) | |
| self.assertEqual(2, self.sp_.PieceToId('</s>')) | |
| self.assertEqual('<unk>', self.sp_.IdToPiece(0)) | |
| self.assertEqual('<s>', self.sp_.IdToPiece(1)) | |
| self.assertEqual('</s>', self.sp_.IdToPiece(2)) | |
| self.assertEqual(0, self.sp_.unk_id()) | |
| self.assertEqual(1, self.sp_.bos_id()) | |
| self.assertEqual(2, self.sp_.eos_id()) | |
| self.assertEqual(-1, self.sp_.pad_id()) | |
| for i in range(self.sp_.GetPieceSize()): | |
| piece = self.sp_.IdToPiece(i) | |
| self.assertEqual(i, self.sp_.PieceToId(piece)) | |
| self.assertEqual(1000, self.sp_.get_piece_size()) | |
| self.assertEqual(0, self.sp_.piece_to_id('<unk>')) | |
| self.assertEqual(1, self.sp_.piece_to_id('<s>')) | |
| self.assertEqual(2, self.sp_.piece_to_id('</s>')) | |
| self.assertEqual('<unk>', self.sp_.id_to_piece(0)) | |
| self.assertEqual('<s>', self.sp_.id_to_piece(1)) | |
| self.assertEqual('</s>', self.sp_.id_to_piece(2)) | |
| for i in range(self.sp_.get_piece_size()): | |
| piece = self.sp_.id_to_piece(i) | |
| self.assertEqual(i, self.sp_.piece_to_id(piece)) | |
| def test_roundtrip(self): | |
| text = 'I saw a girl with a telescope.' | |
| ids = self.sp_.EncodeAsIds(text) | |
| pieces1 = self.sp_.EncodeAsPieces(text) | |
| pieces2 = self.sp_.NBestEncodeAsPieces(text, 10)[0] | |
| self.assertEqual(pieces1, pieces2) | |
| self.assertEqual(text, self.sp_.DecodePieces(pieces1)) | |
| self.assertEqual(text, self.sp_.DecodeIds(ids)) | |
| for n in range(100): | |
| self.assertEqual( | |
| text, | |
| self.sp_.DecodePieces(self.sp_.SampleEncodeAsPieces(text, 64, 0.5)), | |
| ) | |
| self.assertEqual( | |
| text, | |
| self.sp_.DecodePieces(self.sp_.SampleEncodeAsPieces(text, -1, 0.5)), | |
| ) | |
| self.assertEqual( | |
| text, self.sp_.DecodeIds(self.sp_.SampleEncodeAsIds(text, 64, 0.5)) | |
| ) | |
| self.assertEqual( | |
| text, self.sp_.DecodeIds(self.sp_.SampleEncodeAsIds(text, -1, 0.5)) | |
| ) | |
| ids2 = self.sp_.encode_as_ids(text) | |
| pieces3 = self.sp_.encode_as_pieces(text) | |
| pieces4 = self.sp_.nbest_encode_as_pieces(text, 10)[0] | |
| self.assertEqual(pieces3, pieces4) | |
| self.assertEqual(pieces1, pieces3) | |
| self.assertEqual(ids, ids2) | |
| self.assertEqual(text, self.sp_.decode_pieces(pieces3)) | |
| self.assertEqual(text, self.sp_.decode_ids(ids2)) | |
| for n in range(100): | |
| self.assertEqual( | |
| text, | |
| self.sp_.decode_pieces( | |
| self.sp_.sample_encode_as_pieces(text, 64, 0.5) | |
| ), | |
| ) | |
| self.assertEqual( | |
| text, | |
| self.sp_.decode_pieces( | |
| self.sp_.sample_encode_as_pieces(text, -1, 0.5) | |
| ), | |
| ) | |
| self.assertEqual( | |
| text, | |
| self.sp_.decode_ids(self.sp_.sample_encode_as_ids(text, 64, 0.5)), | |
| ) | |
| self.assertEqual( | |
| text, | |
| self.sp_.decode_ids(self.sp_.sample_encode_as_ids(text, -1, 0.5)), | |
| ) | |
| self.assertEqual( | |
| self.sp_.calculate_entropy(text, 0.1), | |
| self.sp_.CalculateEntropy(text, 0.1), | |
| ) | |
| def test_ja_load(self): | |
| self.assertEqual(8000, self.jasp_.GetPieceSize()) | |
| self.assertEqual(0, self.jasp_.PieceToId('<unk>')) | |
| self.assertEqual(1, self.jasp_.PieceToId('<s>')) | |
| self.assertEqual(2, self.jasp_.PieceToId('</s>')) | |
| self.assertEqual('<unk>', self.jasp_.IdToPiece(0)) | |
| self.assertEqual('<s>', self.jasp_.IdToPiece(1)) | |
| self.assertEqual('</s>', self.jasp_.IdToPiece(2)) | |
| for i in range(self.jasp_.GetPieceSize()): | |
| piece = self.jasp_.IdToPiece(i) | |
| self.assertEqual(i, self.jasp_.PieceToId(piece)) | |
| self.assertEqual(8000, self.jasp_.get_piece_size()) | |
| self.assertEqual(0, self.jasp_.piece_to_id('<unk>')) | |
| self.assertEqual(1, self.jasp_.piece_to_id('<s>')) | |
| self.assertEqual(2, self.jasp_.piece_to_id('</s>')) | |
| self.assertEqual('<unk>', self.jasp_.id_to_piece(0)) | |
| self.assertEqual('<s>', self.jasp_.id_to_piece(1)) | |
| self.assertEqual('</s>', self.jasp_.id_to_piece(2)) | |
| for i in range(self.jasp_.get_piece_size()): | |
| piece = self.jasp_.id_to_piece(i) | |
| self.assertEqual(i, self.jasp_.piece_to_id(piece)) | |
| def test_ja_roundtrip(self): | |
| text = '清水寺は京都にある。' | |
| ids = self.jasp_.EncodeAsIds(text) | |
| pieces1 = self.jasp_.EncodeAsPieces(text) | |
| pieces2 = self.jasp_.NBestEncodeAsPieces(text, 10)[0] | |
| self.assertEqual(pieces1, pieces2) | |
| self.assertEqual(text, self.jasp_.DecodePieces(pieces1)) | |
| self.assertEqual(text, self.jasp_.DecodeIds(ids)) | |
| for n in range(100): | |
| self.assertEqual( | |
| text, | |
| self.jasp_.DecodePieces( | |
| self.jasp_.SampleEncodeAsPieces(text, 64, 0.5) | |
| ), | |
| ) | |
| self.assertEqual( | |
| text, | |
| self.jasp_.DecodePieces( | |
| self.jasp_.SampleEncodeAsPieces(text, -1, 0.5) | |
| ), | |
| ) | |
| ids2 = self.jasp_.encode_as_ids(text) | |
| pieces3 = self.jasp_.encode_as_pieces(text) | |
| pieces4 = self.jasp_.nbest_encode_as_pieces(text, 10)[0] | |
| self.assertEqual(pieces3, pieces4) | |
| self.assertEqual(pieces1, pieces3) | |
| self.assertEqual(ids, ids2) | |
| self.assertEqual(text, self.jasp_.decode_pieces(pieces1)) | |
| self.assertEqual(text, self.jasp_.decode_ids(ids2)) | |
| for n in range(100): | |
| self.assertEqual( | |
| text, | |
| self.jasp_.decode_pieces( | |
| self.jasp_.sample_encode_as_pieces(text, 64, 0.5) | |
| ), | |
| ) | |
| self.assertEqual( | |
| text, | |
| self.jasp_.decode_pieces( | |
| self.jasp_.sample_encode_as_pieces(text, -1, 0.5) | |
| ), | |
| ) | |
| self.assertEqual( | |
| self.jasp_.calculate_entropy(text, 0.1), | |
| self.jasp_.CalculateEntropy(text, 0.1), | |
| ) | |
| def test_train(self): | |
| spm.SentencePieceTrainer.Train( | |
| '--input=' | |
| + os.path.join(data_dir, 'botchan.txt') | |
| + ' --model_prefix=m --vocab_size=1000' | |
| ) | |
| sp = spm.SentencePieceProcessor() | |
| sp.Load('m.model') | |
| with open(os.path.join(data_dir, 'botchan.txt'), 'r') as file: | |
| for line in file: | |
| sp.DecodePieces(sp.EncodeAsPieces(line)) | |
| sp.DecodeIds(sp.EncodeAsIds(line)) | |
| def test_train_iterator(self): | |
| spm.SentencePieceTrainer.Train( | |
| '--input=' | |
| + os.path.join(data_dir, 'botchan.txt') | |
| + ' --model_prefix=m --vocab_size=1000' | |
| ) | |
| # Load as 'rb' for Python3.5/2.7. | |
| os1 = io.BytesIO() | |
| os2 = io.BytesIO() | |
| # suppress logging (redirect to /dev/null) | |
| spm.SentencePieceTrainer.train( | |
| input=os.path.join(data_dir, 'botchan.txt'), | |
| model_prefix='m', | |
| vocab_size=1000, | |
| logstream=open(os.devnull, 'w'), | |
| ) | |
| with open(os.path.join(data_dir, 'botchan.txt'), 'rb') as is1: | |
| spm.SentencePieceTrainer.train( | |
| sentence_iterator=is1, | |
| model_prefix='m', | |
| vocab_size=1000, | |
| logstream=open(os.devnull, 'w'), | |
| ) | |
| spm.SentencePieceTrainer.train( | |
| input=os.path.join(data_dir, 'botchan.txt'), | |
| model_writer=os1, | |
| vocab_size=1000, | |
| logstream=open(os.devnull, 'w'), | |
| ) | |
| with open(os.path.join(data_dir, 'botchan.txt'), 'rb') as is2: | |
| spm.SentencePieceTrainer.train( | |
| sentence_iterator=is2, | |
| model_writer=os2, | |
| vocab_size=1000, | |
| logstream=open(os.devnull, 'w'), | |
| ) | |
| sp1 = spm.SentencePieceProcessor(model_proto=os1.getvalue()) | |
| sp2 = spm.SentencePieceProcessor(model_proto=os2.getvalue()) | |
| self.assertEqual( | |
| [sp1.id_to_piece(i) for i in range(sp1.get_piece_size())], | |
| [sp2.id_to_piece(i) for i in range(sp2.get_piece_size())], | |
| ) | |
| def test_train_kwargs(self): | |
| # suppress logging (redirect to /dev/null) | |
| spm.SentencePieceTrainer.train( | |
| input=[os.path.join(data_dir, 'botchan.txt')], | |
| model_prefix='m', | |
| vocab_size=1002, | |
| user_defined_symbols=['foo', 'bar', ',', ' ', '\t', '\b', '\n', '\r'], | |
| logstream=open(os.devnull, 'w'), | |
| ) | |
| sp = spm.SentencePieceProcessor() | |
| sp.Load('m.model') | |
| with open(os.path.join(data_dir, 'botchan.txt'), 'r') as file: | |
| for line in file: | |
| sp.DecodePieces(sp.EncodeAsPieces(line)) | |
| sp.DecodeIds(sp.EncodeAsIds(line)) | |
| s = 'hello\tworld\r\nthis\tis a \b pen' | |
| self.assertEqual(s, sp.decode(sp.encode(s))) | |
| def test_serialized_proto(self): | |
| text = 'I saw a girl with a telescope.' | |
| s1 = self.sp_.EncodeAsSerializedProto(text) | |
| s2 = self.sp_.SampleEncodeAsSerializedProto(text, 10, 0.2) | |
| s3 = self.sp_.NBestEncodeAsSerializedProto(text, 10) | |
| s4 = self.sp_.DecodePiecesAsSerializedProto(['foo', 'bar']) | |
| s5 = self.sp_.DecodeIdsAsSerializedProto([20, 30]) | |
| t1 = self.sp_.encode_as_serialized_proto(text) | |
| t2 = self.sp_.sample_encode_as_serialized_proto(text, 10, 0.2) | |
| t3 = self.sp_.nbest_encode_as_serialized_proto(text, 10) | |
| t4 = self.sp_.decode_pieces_as_serialized_proto(['foo', 'bar']) | |
| t5 = self.sp_.decode_ids_as_serialized_proto([20, 30]) | |
| y1 = self.sp_.encode(text, out_type='serialized_proto') | |
| y2 = self.sp_.encode( | |
| text, enable_sampling=True, out_type='serialized_proto' | |
| ) | |
| y3 = self.sp_.nbest_encode(text, out_type='serialized_proto', nbest_size=10) | |
| y4 = self.sp_.decode(['foo', 'bar'], out_type='serialized_proto') | |
| y5 = self.sp_.decode([20, 30], out_type='serialized_proto') | |
| self.assertEqual(type(s1), bytes) | |
| self.assertEqual(type(s2), bytes) | |
| self.assertEqual(type(t2), bytes) | |
| self.assertEqual(type(s3), bytes) | |
| self.assertEqual(type(s4), bytes) | |
| self.assertEqual(type(s5), bytes) | |
| self.assertEqual(s1, t1) | |
| self.assertEqual(s3, t3) | |
| self.assertEqual(s4, t4) | |
| self.assertEqual(s5, t5) | |
| self.assertEqual(s1, y1) | |
| self.assertEqual(s3, y3) | |
| self.assertEqual(s4, y4) | |
| self.assertEqual(s5, y5) | |
| ids = self.jasp_.EncodeAsIds(text) | |
| pieces = self.jasp_.EncodeAsPieces(text) | |
| s1 = self.jasp_.EncodeAsSerializedProto(text) | |
| s2 = self.jasp_.DecodeIdsAsSerializedProto(ids) | |
| s3 = self.jasp_.DecodePiecesAsSerializedProto(ids) | |
| self.assertEqual(s2, s1) | |
| self.assertEqual(s3, s1) | |
| def test_decode_bytes(self): | |
| texts = ['Hello world', '清水寺は京都にある。'] | |
| ids = self.jasp_.encode(texts, out_type=int) | |
| s1 = self.jasp_.decode(ids, out_type=bytes) | |
| s2 = self.jasp_.decode(ids, out_type=str) | |
| self.assertEqual(len(s1), 2) | |
| self.assertEqual(type(s1[0]), bytes) | |
| self.assertEqual(type(s1[1]), bytes) | |
| self.assertEqual(len(s2), 2) | |
| self.assertEqual(type(s2[0]), str) | |
| self.assertEqual(type(s2[1]), str) | |
| self.assertEqual(s1[0].decode(encoding='utf-8'), s2[0]) | |
| self.assertEqual(s1[1].decode(encoding='utf-8'), s2[1]) | |
| text = 'Hello world' | |
| ids = self.jasp_.encode(text, out_type=int) | |
| s1 = self.jasp_.decode(ids, out_type=bytes) | |
| s2 = self.jasp_.decode(ids, out_type=str) | |
| self.assertEqual(type(s1), bytes) | |
| self.assertEqual(type(s2), str) | |
| self.assertEqual(s1.decode(encoding='utf-8'), s2) | |
| x = self.jasp_.encode(text, out_type='immutable_proto') | |
| self.assertEqual(x.text, x.text_as_bytes.decode(encoding='utf-8')) | |
| for sp in x.pieces: | |
| self.assertEqual(sp.piece, sp.piece_as_bytes.decode(encoding='utf-8')) | |
| self.assertEqual(sp.surface, sp.surface_as_bytes.decode(encoding='utf-8')) | |
| x = self.jasp_.decode(ids, out_type='immutable_proto') | |
| self.assertEqual(x.text, x.text_as_bytes.decode(encoding='utf-8')) | |
| for sp in x.pieces: | |
| self.assertEqual(sp.piece, sp.piece_as_bytes.decode(encoding='utf-8')) | |
| self.assertEqual(sp.surface, sp.surface_as_bytes.decode(encoding='utf-8')) | |
| def test_immutable_proto(self): | |
| text = 'I saw a girl with a telescope.' | |
| s1 = self.sp_.EncodeAsImmutableProto(text) | |
| s2 = self.sp_.SampleEncodeAsImmutableProto(text, 10, 0.2) | |
| s3 = self.sp_.NBestEncodeAsImmutableProto(text, 10) | |
| s4 = self.sp_.DecodePiecesAsImmutableProto(['foo', 'bar']) | |
| s5 = self.sp_.DecodeIdsAsImmutableProto([20, 30]) | |
| print(s1) | |
| print(s2) | |
| print(s3) | |
| print(s4) | |
| print(s5) | |
| t1 = self.sp_.encode_as_immutable_proto(text) | |
| t2 = self.sp_.sample_encode_as_immutable_proto(text, 10, 0.2) | |
| t3 = self.sp_.nbest_encode_as_immutable_proto(text, 10) | |
| t4 = self.sp_.decode_pieces_as_immutable_proto(['foo', 'bar']) | |
| t5 = self.sp_.decode_ids_as_immutable_proto([20, 30]) | |
| y1 = self.sp_.encode(text, out_type='immutable_proto') | |
| y2 = self.sp_.encode(text, enable_sampling=True, out_type='immutable_proto') | |
| y3 = self.sp_.nbest_encode(text, out_type='immutable_proto', nbest_size=10) | |
| y4 = self.sp_.decode(['foo', 'bar'], out_type='immutable_proto') | |
| y5 = self.sp_.decode([20, 30], out_type='immutable_proto') | |
| self.assertEqual(s1, t1) | |
| self.assertEqual(s3, t3) | |
| self.assertEqual(s4, t4) | |
| self.assertEqual(s5, t5) | |
| self.assertEqual(s1, y1) | |
| self.assertEqual(s3, y3) | |
| self.assertEqual(s4, y4) | |
| self.assertEqual(s5, y5) | |
| hset_piece = defaultdict(int) | |
| # eq test | |
| for i in range(len(s1.pieces)): | |
| self.assertEqual(s1.pieces[i], t1.pieces[i]) | |
| hset_piece[s1.pieces[i]] += 1 | |
| hset_piece[t1.pieces[i]] += 1 | |
| self.assertEqual(len(hset_piece), len(s1.pieces)) | |
| # has test | |
| hset = defaultdict(int) | |
| hset[s1] += 1 | |
| hset[t1] += 1 | |
| hset[s3] += 1 | |
| hset[t3] += 1 | |
| self.assertEqual(len(hset), 2) | |
| self.assertEqual(hset[s1], 2) | |
| self.assertEqual(hset[s3], 2) | |
| self.assertEqual(hset[t1], 2) | |
| self.assertEqual(hset[t3], 2) | |
| x1 = self.sp_.encode_as_serialized_proto(text) | |
| x2 = self.sp_.sample_encode_as_serialized_proto(text, 10, 0.2) | |
| x3 = self.sp_.nbest_encode_as_serialized_proto(text, 10) | |
| x4 = self.sp_.decode_pieces_as_serialized_proto(['foo', 'bar']) | |
| x5 = self.sp_.decode_ids_as_serialized_proto([20, 30]) | |
| self.assertEqual(x1, t1.SerializeAsString()) | |
| self.assertEqual(x3, t3.SerializeAsString()) | |
| self.assertEqual(x4, t4.SerializeAsString()) | |
| self.assertEqual(x5, t5.SerializeAsString()) | |
| v1 = self.sp_.EncodeAsIds(text) | |
| v2 = self.sp_.EncodeAsPieces(text) | |
| self.assertEqual([x.id for x in s1.pieces], v1) | |
| self.assertEqual([x.piece for x in s1.pieces], v2) | |
| self.assertEqual(text, s1.text) | |
| surfaces1 = [s1.text[x.begin : x.end] for x in s1.pieces] | |
| surfaces2 = [x.surface for x in s1.pieces] | |
| self.assertEqual(surfaces1, surfaces2) | |
| ids = [] | |
| for i in range(len(s1.pieces)): | |
| ids.append(s1.pieces[i].id) | |
| self.assertEqual(ids, v1) | |
| pieces = [] | |
| for i in range(len(s1.pieces)): | |
| pieces.append(s1.pieces[i].piece) | |
| self.assertEqual(pieces, v2) | |
| for v in s3.nbests: | |
| self.assertEqual(text, v.text) | |
| self.assertEqual(self.sp_.Decode([x.id for x in v.pieces]), text) | |
| for i in range(len(s3.nbests)): | |
| self.assertEqual(text, s3.nbests[i].text) | |
| self.assertEqual( | |
| self.sp_.Decode([x.id for x in s3.nbests[i].pieces]), text | |
| ) | |
| # slice | |
| self.assertEqual(s1.pieces[::-1], list(reversed(s1.pieces))) | |
| self.assertEqual(s3.nbests[::-1], list(reversed(s3.nbests))) | |
| # Japanese offset | |
| s1 = self.jasp_.EncodeAsImmutableProto( | |
| '吾輩は猫である。Hello world. ABC 123' | |
| ) | |
| surfaces1 = [s1.text[x.begin : x.end] for x in s1.pieces] | |
| surfaces2 = [x.surface for x in s1.pieces] | |
| self.assertEqual(surfaces1, surfaces2) | |
| ids = [x.id for x in s1.pieces] | |
| s2 = self.jasp_.DecodeIdsAsImmutableProto(ids) | |
| self.assertEqual(s2, s1) | |
| pieces = [x.piece for x in s1.pieces] | |
| s2 = self.jasp_.DecodePiecesAsImmutableProto(pieces) | |
| self.assertEqual(s2, s1) | |
| def test_new_api(self): | |
| sp = spm.SentencePieceProcessor( | |
| model_file=os.path.join('test', 'test_model.model') | |
| ) | |
| text = 'hello world' | |
| text2 = 'Tokyo' | |
| ids = self.sp_.EncodeAsIds(text) | |
| ids2 = self.sp_.EncodeAsIds(text2) | |
| pieces = self.sp_.EncodeAsPieces(text) | |
| pieces2 = self.sp_.EncodeAsPieces(text2) | |
| sprotos = self.sp_.EncodeAsSerializedProto(text) | |
| sproto2 = self.sp_.EncodeAsSerializedProto(text2) | |
| iprotos = self.sp_.EncodeAsImmutableProto(text) | |
| iprotos2 = self.sp_.EncodeAsImmutableProto(text2) | |
| self.assertEqual(sp.encode(text, out_type=int), ids) | |
| self.assertEqual(sp.encode(text, out_type=str), pieces) | |
| self.assertEqual(sp.encode(text, out_type='serialized_proto'), sprotos) | |
| self.assertEqual(sp.encode(text, out_type='immutable_proto'), iprotos) | |
| self.assertEqual(sp.encode([text], out_type=int), [ids]) | |
| self.assertEqual(sp.encode([text], out_type=str), [pieces]) | |
| self.assertEqual(sp.encode([text], out_type='serialized_proto'), [sprotos]) | |
| self.assertEqual(sp.encode([text], out_type='immutable_proto'), [iprotos]) | |
| self.assertEqual(len(iprotos.pieces), len(pieces)) | |
| self.assertEqual(len(iprotos.pieces), len(ids)) | |
| self.assertEqual(iprotos.text, text) | |
| self.assertEqual(len(iprotos2.pieces), len(pieces2)) | |
| self.assertEqual(len(iprotos2.pieces), len(ids2)) | |
| self.assertEqual(iprotos2.text, text2) | |
| for i in range(len(iprotos.pieces)): | |
| self.assertEqual(ids[i], iprotos.pieces[i].id) | |
| self.assertEqual(pieces[i], iprotos.pieces[i].piece) | |
| for i, piece in enumerate(iprotos.pieces): | |
| self.assertEqual(ids[i], piece.id) | |
| self.assertEqual(pieces[i], piece.piece) | |
| for i in range(len(iprotos2.pieces)): | |
| self.assertEqual(ids2[i], iprotos2.pieces[i].id) | |
| self.assertEqual(pieces2[i], iprotos2.pieces[i].piece) | |
| for i, piece in enumerate(iprotos2.pieces): | |
| self.assertEqual(ids2[i], piece.id) | |
| self.assertEqual(pieces2[i], piece.piece) | |
| detok_ids = self.sp_.DecodeIds(ids) | |
| detok_pieces = self.sp_.DecodePieces(pieces) | |
| self.assertEqual(sp.decode(ids), detok_ids) | |
| self.assertEqual(sp.decode(pieces), detok_pieces) | |
| self.assertEqual(sp.decode([]), '') | |
| self.assertEqual(sp.decode([[]]), ['']) | |
| # add_bos, add_eos, reverse | |
| self.assertEqual([sp.bos_id()] + ids, sp.encode(text, add_bos=True)) | |
| self.assertEqual(ids + [sp.eos_id()], sp.encode(text, add_eos=True)) | |
| self.assertEqual(ids + [sp.eos_id()], sp.EncodeAsIds(text, add_eos=True)) | |
| rids = ids[:] | |
| rids.reverse() | |
| self.assertEqual(rids, sp.encode(text, reverse=True)) | |
| self.assertEqual(rids, sp.EncodeAsIds(text, reverse=True)) | |
| # different shape. | |
| self.assertEqual([ids, ids2], sp.encode([text, text2])) | |
| self.assertEqual([pieces, pieces2], sp.encode([text, text2], out_type=str)) | |
| self.assertEqual([text, text2], sp.decode([ids, ids2])) | |
| self.assertEqual([text, text2], sp.decode([pieces, pieces2])) | |
| pieces = list(reversed(self.sp_.EncodeAsPieces(text))) | |
| self.assertEqual(pieces, sp.encode(text, reverse=True, out_type=str)) | |
| # emit unk piece | |
| unk_char = '藤' | |
| pieces = self.sp_.EncodeAsIds(unk_char, emit_unk_piece=True) | |
| pieces2 = self.sp_.encode(unk_char, out_type=int, emit_unk_piece=True) | |
| self.assertEqual(pieces[1], sp.unk_id()) | |
| self.assertEqual(pieces2[1], sp.unk_id()) | |
| self.assertEqual(pieces, pieces2) | |
| pieces = self.sp_.EncodeAsPieces(unk_char, emit_unk_piece=True) | |
| pieces2 = self.sp_.encode(unk_char, out_type=str, emit_unk_piece=True) | |
| self.assertEqual(pieces[1], '<unk>') | |
| self.assertEqual(pieces2[1], '<unk>') | |
| self.assertEqual(pieces, pieces2) | |
| pieces = self.sp_.EncodeAsPieces(unk_char, emit_unk_piece=False) | |
| pieces2 = self.sp_.encode(unk_char, out_type=str, emit_unk_piece=False) | |
| self.assertEqual(pieces[1], unk_char) | |
| self.assertEqual(pieces2[1], unk_char) | |
| self.assertEqual(pieces, pieces2) | |
| def test_new_api_init(self): | |
| sp = spm.SentencePieceProcessor( | |
| model_file=os.path.join('test', 'test_model.model'), | |
| add_bos=True, | |
| add_eos=True, | |
| out_type=str, | |
| ) | |
| text = 'hello world' | |
| pieces = ['<s>'] + self.sp_.EncodeAsPieces(text) + ['</s>'] | |
| self.assertEqual(pieces, sp.encode(text)) | |
| pieces = self.sp_.EncodeAsPieces(text) + ['</s>'] | |
| self.assertEqual(pieces, sp.encode(text, add_bos=False, add_eos=True)) | |
| def test_sampling(self): | |
| sp = self.sp_ | |
| for out_type in [str, int, 'serialized_proto', 'immutable_proto']: | |
| ids = defaultdict(int) | |
| for n in range(100): | |
| out = sp.encode('hello world', out_type=out_type, enable_sampling=True) | |
| if type(out) is list: | |
| out = tuple(out) | |
| ++ids[out] | |
| self.assertGreater(len(ids), 1) | |
| ids2 = defaultdict(int) | |
| for n in range(100): | |
| out = sp.encode('hello world', out_type=out_type, enable_sampling=False) | |
| if type(out) is list: | |
| out = tuple(out) | |
| ++ids2[out] | |
| self.assertEqual(len(ids2), 1) | |
| out = sp.encode( | |
| ['hello world', 'this is a test'], | |
| out_type=out_type, | |
| enable_sampling=True, | |
| ) | |
| self.assertEqual(len(out), 2) | |
| out = sp.encode( | |
| ['hello world', 'this is a test'], | |
| out_type=out_type, | |
| enable_sampling=False, | |
| ) | |
| self.assertEqual(len(out), 2) | |
| def test_nbest(self): | |
| sp = self.sp_ | |
| text = 'hello world' | |
| text2 = 'I have a pen.' | |
| for out_type in [str, int, 'serialized_proto', 'immutable_proto']: | |
| results = sp.nbest_encode(text, nbest_size=10, out_type=out_type) | |
| self.assertEqual( | |
| results, sp.NBestEncode(text, nbest_size=10, out_type=out_type) | |
| ) | |
| if out_type in [str, int]: | |
| for n in results: | |
| self.assertEqual(sp.decode(n), text) | |
| for n in sp.decode(results): | |
| self.assertEqual(n, text) | |
| # batch test | |
| results = sp.nbest_encode([text, text2], nbest_size=10, out_type=out_type) | |
| self.assertEqual( | |
| results, | |
| sp.NBestEncode([text, text2], nbest_size=10, out_type=out_type), | |
| ) | |
| self.assertEqual(len(results), 2) | |
| if out_type in [str, int]: | |
| for n in results[0]: | |
| self.assertEqual(sp.decode(n), text) | |
| for n in results[1]: | |
| self.assertEqual(sp.decode(n), text2) | |
| decoded = sp.decode(results[0]) | |
| self.assertEqual(len(decoded), 10) | |
| for n in decoded: | |
| self.assertEqual(n, text) | |
| decoded = sp.decode(results[1]) | |
| self.assertEqual(len(decoded), 10) | |
| for n in decoded: | |
| self.assertEqual(n, text2) | |
| self.assertEqual( | |
| sp.nbest_encode(text, nbest_size=10, out_type=str), | |
| sp.nbest_encode_as_pieces(text, nbest_size=10), | |
| ) | |
| self.assertEqual( | |
| sp.nbest_encode(text, nbest_size=10, out_type=int), | |
| sp.nbest_encode_as_ids(text, nbest_size=10), | |
| ) | |
| self.assertEqual( | |
| sp.nbest_encode(text, nbest_size=10, out_type='serialized_proto'), | |
| sp.nbest_encode_as_serialized_proto(text, nbest_size=10), | |
| ) | |
| self.assertEqual( | |
| sp.nbest_encode(text, nbest_size=10, out_type='immutable_proto'), | |
| sp.nbest_encode_as_immutable_proto(text, nbest_size=10), | |
| ) | |
| def test_sample_and_score(self): | |
| sp = self.sp_ | |
| text = 'hello world' | |
| text2 = 'I have a pen.' | |
| for out_type in [str, int, 'serialized_proto', 'immutable_proto']: | |
| results = sp.sample_encode_and_score( | |
| text, wor=True, num_samples=10, out_type=out_type | |
| ) | |
| results = sp.SampleEncodeAndScore( | |
| text, wor=False, num_samples=10, out_type=out_type | |
| ) | |
| if out_type in [str, int]: | |
| for n in results: | |
| self.assertEqual(sp.decode(n[0]), text) | |
| results = sp.sample_encode_and_score( | |
| [text, text2], wor=True, num_samples=10, out_type=out_type | |
| ) | |
| results = sp.SampleEncodeAndScore( | |
| [text, text2], wor=True, num_samples=10, out_type=out_type | |
| ) | |
| if out_type in [str, int]: | |
| for n in results[0]: | |
| self.assertEqual(sp.decode(n[0]), text) | |
| for n in results[1]: | |
| self.assertEqual(sp.decode(n[0]), text2) | |
| sp.sample_encode_and_score_as_pieces(text, 10) | |
| sp.sample_encode_and_score_as_ids(text, 10) | |
| sp.sample_encode_and_score_as_immutable_proto(text, 10) | |
| sp.sample_encode_and_score_as_serialized_proto(text, 10) | |
| def test_valid_range(self): | |
| size = self.sp_.piece_size() | |
| funcs = [ | |
| 'IdToPiece', | |
| 'GetScore', | |
| 'IsUnknown', | |
| 'IsControl', | |
| 'IsUnused', | |
| 'IsByte', | |
| 'DecodeIds', | |
| 'DecodeIdsAsSerializedProto', | |
| ] | |
| for m in funcs: | |
| getattr(self.sp_, m)([10, 20, 30]) | |
| for m in funcs: | |
| try: | |
| getattr(self.sp_, m)([size]) | |
| self.assertTrue(False) | |
| except: | |
| self.assertTrue(True) | |
| def test_batch(self): | |
| sp = spm.SentencePieceProcessor( | |
| model_file=os.path.join('test', 'test_model.model') | |
| ) | |
| with open(os.path.join(data_dir, 'botchan.txt'), 'r') as file: | |
| texts = file.readlines() | |
| for out_type in [str, int, 'serialized_proto', 'immutable_proto']: | |
| r1 = sp.encode(texts, out_type=out_type, num_threads=None) | |
| r2 = sp.encode(texts, out_type=out_type, num_threads=1) | |
| r3 = sp.encode(texts, out_type=out_type, num_threads=-1) | |
| r4 = sp.encode(texts, out_type=out_type, num_threads=8) | |
| r5 = [sp.encode(s, out_type=out_type) for s in texts] | |
| self.assertEqual(r1, r2) | |
| self.assertEqual(r1, r3) | |
| self.assertEqual(r1, r4) | |
| self.assertEqual(r1, r5) | |
| if out_type in [str, int]: | |
| d1 = sp.decode(r1, num_threads=None) | |
| d2 = sp.decode(r2, num_threads=1) | |
| d3 = sp.decode(r3, num_threads=-1) | |
| d4 = sp.decode(r4, num_threads=8) | |
| d5 = [sp.decode(s) for s in r5] | |
| self.assertEqual(d1, d2) | |
| self.assertEqual(d1, d3) | |
| self.assertEqual(d1, d4) | |
| self.assertEqual(d1, d5) | |
| e1 = sp.calculate_entropy(texts, alpha=1.0, num_threads=10) | |
| e2 = sp.CalculateEntropy(texts, alpha=1.0, num_threads=10) | |
| e3 = [sp.calculate_entropy(s, alpha=1.0) for s in texts] | |
| self.assertEqual(e1, e2) | |
| self.assertEqual(e1, e3) | |
| def test_pickle(self): | |
| with open('sp.pickle', 'wb') as f: | |
| pickle.dump(self.sp_, f) | |
| id1 = self.sp_.encode('hello world.', out_type=int) | |
| with open('sp.pickle', 'rb') as f: | |
| sp = pickle.load(f) | |
| id2 = sp.encode('hello world.', out_type=int) | |
| self.assertEqual(id1, id2) | |
| def test_global_params(self): | |
| spm.SetRandomGeneratorSeed(0) | |
| spm.SetMinLogLevel(2) | |
| spm.set_random_generator_seed(1) | |
| spm.set_min_log_level(3) | |
| def test_normalize(self): | |
| sp = spm.SentencePieceProcessor( | |
| model_file=os.path.join('test', 'test_model.model') | |
| ) | |
| self.assertEqual('▁KADOKAWAABC', sp.normalize('ＫＡＤＯＫＡＷＡABC')) | |
| self.assertEqual('▁KADOKAWAABC', sp.Normalize('ＫＡＤＯＫＡＷＡABC')) | |
| x = sp.Normalize('ＫＡＤＯＫＡＷＡABC', with_offsets=True) | |
| self.assertEqual('▁KADOKAWAABC', x[0]) | |
| self.assertEqual([0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], x[1]) | |
| x = sp.Normalize('ＫＡＤＯＫＡＷＡABC'.encode('utf8'), with_offsets=True) | |
| self.assertEqual('▁KADOKAWAABC'.encode('utf8'), x[0]) | |
| self.assertEqual( | |
| [0, 0, 0, 0, 3, 6, 9, 12, 15, 18, 21, 24, 25, 26, 27], x[1] | |
| ) | |
| self.assertEqual( | |
| ['▁KADOKAWAABC', '▁平成'], sp.normalize(['ＫＡＤＯＫＡＷＡABC', '㍻']) | |
| ) | |
| self.assertEqual( | |
| ['▁KADOKAWAABC', '▁平成'], sp.Normalize(['ＫＡＤＯＫＡＷＡABC', '㍻']) | |
| ) | |
| x = sp.Normalize( | |
| ['ＫＡＤＯＫＡＷＡABC'.encode('utf8'), '㍻'.encode('utf8')], | |
| with_offsets=True, | |
| ) | |
| self.assertEqual(len(x), 2) | |
| self.assertEqual('▁KADOKAWAABC'.encode('utf8'), x[0][0]) | |
| self.assertEqual( | |
| [0, 0, 0, 0, 3, 6, 9, 12, 15, 18, 21, 24, 25, 26, 27], x[0][1] | |
| ) | |
| x = sp.Normalize(['ＫＡＤＯＫＡＷＡABC', '㍻'], with_offsets=True) | |
| self.assertEqual(len(x), 2) | |
| self.assertEqual('▁KADOKAWAABC', x[0][0]) | |
| self.assertEqual([0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], x[0][1]) | |
| self.assertEqual('▁平成', x[1][0]) | |
| self.assertEqual([0, 0, 0, 1], x[1][1]) | |
| def test_normalizer(self): | |
| sp = spm.SentencePieceNormalizer( | |
| model_file=os.path.join('test', 'test_model.model') | |
| ) | |
| self.assertEqual('KADOKAWAABC', sp.normalize('ＫＡＤＯＫＡＷＡABC')) | |
| self.assertEqual('KADOKAWAABC', sp.Normalize('ＫＡＤＯＫＡＷＡABC')) | |
| x = sp.Normalize('ＫＡＤＯＫＡＷＡABC'.encode('utf8'), with_offsets=True) | |
| self.assertEqual('KADOKAWAABC'.encode('utf8'), x[0]) | |
| self.assertEqual([0, 3, 6, 9, 12, 15, 18, 21, 24, 25, 26, 27], x[1]) | |
| x = sp.Normalize('ＫＡＤＯＫＡＷＡABC', with_offsets=True) | |
| self.assertEqual('KADOKAWAABC', x[0]) | |
| self.assertEqual([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], x[1]) | |
| self.assertEqual( | |
| ['KADOKAWAABC', '平成'], sp.normalize(['ＫＡＤＯＫＡＷＡABC', '㍻']) | |
| ) | |
| self.assertEqual( | |
| ['KADOKAWAABC', '平成'], sp.Normalize(['ＫＡＤＯＫＡＷＡABC', '㍻']) | |
| ) | |
| x = sp.Normalize( | |
| ['ＫＡＤＯＫＡＷＡABC'.encode('utf8'), '㍻'.encode('utf8')], | |
| with_offsets=True, | |
| ) | |
| self.assertEqual(len(x), 2) | |
| self.assertEqual('KADOKAWAABC'.encode('utf8'), x[0][0]) | |
| self.assertEqual([0, 3, 6, 9, 12, 15, 18, 21, 24, 25, 26, 27], x[0][1]) | |
| x = sp.Normalize(['ＫＡＤＯＫＡＷＡABC', '㍻'], with_offsets=True) | |
| self.assertEqual(len(x), 2) | |
| self.assertEqual('KADOKAWAABC', x[0][0]) | |
| self.assertEqual([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], x[0][1]) | |
| self.assertEqual('平成', x[1][0]) | |
| self.assertEqual([0, 0, 1], x[1][1]) | |
| sp = spm.SentencePieceNormalizer( | |
| model_file=os.path.join('test', 'test_model.model'), | |
| add_dummy_prefix=True, | |
| escape_whitespaces=True, | |
| remove_extra_whitespaces=False, | |
| ) | |
| self.assertEqual('▁hello▁▁world', sp.normalize('hello world')) | |
| sp = spm.SentencePieceNormalizer( | |
| model_file=os.path.join('test', 'test_model.model'), | |
| add_dummy_prefix=True, | |
| escape_whitespaces=True, | |
| remove_extra_whitespaces=True, | |
| ) | |
| self.assertEqual('▁hello▁world', sp.normalize(' hello world ')) | |
| sp = spm.SentencePieceNormalizer( | |
| model_file=os.path.join('test', 'test_model.model'), | |
| add_dummy_prefix=False, | |
| escape_whitespaces=False, | |
| remove_extra_whitespaces=True, | |
| ) | |
| self.assertEqual('hello world', sp.normalize(' hello world ')) | |
| def test_normalizer_rule(self): | |
| sp = spm.SentencePieceNormalizer(rule_name='identity') | |
| self.assertEqual('ＡＢＣ', sp.Normalize('ＡＢＣ')) | |
| sp = spm.SentencePieceNormalizer(rule_name='nfkc_cf') | |
| self.assertEqual('abc', sp.Normalize('ＡＢＣ')) | |
| def test_override_normalize_spec(self): | |
| sp = spm.SentencePieceProcessor( | |
| model_file=os.path.join('test', 'test_model.model') | |
| ) | |
| self.assertEqual( | |
| sp.EncodeAsPieces(' hello world '), ['▁he', 'll', 'o', '▁world'] | |
| ) | |
| sp.override_normalizer_spec(add_dummy_prefix=False) | |
| sp.override_normalizer_spec(remove_extra_whitespaces=False) | |
| sp.override_normalizer_spec(escape_whitespaces=False) | |
| self.assertEqual( | |
| sp.EncodeAsPieces(' hello world '), | |
| [' ', 'he', 'll', 'o', ' ', 'w', 'or', 'l', 'd', ' '], | |
| ) | |
| def suite(): | |
| suite = unittest.TestSuite() | |
| suite.addTests(unittest.makeSuite(TestSentencepieceProcessor)) | |
| return suite | |
| if __name__ == '__main__': | |
| unittest.main() | |