import pytest import os from fastNLP.embeddings.torch import StaticEmbedding from fastNLP import Vocabulary from fastNLP.envs.imports import _NEED_IMPORT_TORCH if _NEED_IMPORT_TORCH: import torch import numpy as np tests_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')) @pytest.mark.torch class TestLoad: def test_norm1(self): # 测试只对可以找到的norm vocab = Vocabulary().add_word_lst(['the', 'a', 'notinfile']) embed = StaticEmbedding(vocab, model_dir_or_name=tests_folder+'/helpers/data/embedding/small_static_embedding/' 'glove.6B.50d_test.txt', only_norm_found_vector=True) assert round(torch.norm(embed(torch.LongTensor([[2]]))).item(), 4) == 1 assert torch.norm(embed(torch.LongTensor([[4]]))).item() != 1 def test_norm2(self): # 测试对所有都norm vocab = Vocabulary().add_word_lst(['the', 'a', 'notinfile']) embed = StaticEmbedding(vocab, model_dir_or_name=tests_folder+'/helpers/data/embedding/small_static_embedding/' 'glove.6B.50d_test.txt', normalize=True) assert round(torch.norm(embed(torch.LongTensor([[2]]))).item(), 4) == 1 assert round(torch.norm(embed(torch.LongTensor([[4]]))).item(), 4) == 1 def test_dropword(self): # 测试是否可以通过drop word vocab = Vocabulary().add_word_lst([chr(i) for i in range(1, 200)]) embed = StaticEmbedding(vocab, model_dir_or_name=None, embedding_dim=10, dropout=0.1, word_dropout=0.4) for i in range(10): length = torch.randint(1, 50, (1,)).item() batch = torch.randint(1, 4, (1,)).item() words = torch.randint(1, 200, (batch, length)).long() embed(words) def test_only_use_pretrain_word(self): def check_word_unk(words, vocab, embed): for word in words: assert embed(torch.LongTensor([vocab.to_index(word)])).tolist()[0] == embed(torch.LongTensor([1])).tolist()[0] def check_vector_equal(words, vocab, embed, embed_dict, lower=False): for word in words: index = vocab.to_index(word) v1 = embed(torch.LongTensor([index])).tolist()[0] if lower: word = word.lower() v2 = embed_dict[word] for v1i, v2i in zip(v1, v2): assert np.allclose(v1i, v2i) embed_dict = read_static_embed(tests_folder+'/helpers/data/embedding/small_static_embedding/' 'glove.6B.50d_test.txt') # 测试是否只使用pretrain的word vocab = Vocabulary().add_word_lst(['the', 'a', 'notinfile']) vocab.add_word('of', no_create_entry=True) embed = StaticEmbedding(vocab, model_dir_or_name=tests_folder+'/helpers/data/embedding/small_static_embedding/' 'glove.6B.50d_test.txt', only_use_pretrain_word=True) # notinfile应该被置为unk check_vector_equal(['the', 'a', 'of'], vocab, embed, embed_dict) check_word_unk(['notinfile'], vocab, embed) # 测试在大小写情况下的使用 vocab = Vocabulary().add_word_lst(['The', 'a', 'notinfile']) vocab.add_word('Of', no_create_entry=True) embed = StaticEmbedding(vocab, model_dir_or_name=tests_folder+'/helpers/data/embedding/small_static_embedding/' 'glove.6B.50d_test.txt', only_use_pretrain_word=True) check_word_unk(['The', 'Of', 'notinfile'], vocab, embed) # 这些词应该找不到 check_vector_equal(['a'], vocab, embed, embed_dict) embed = StaticEmbedding(vocab, model_dir_or_name=tests_folder+'/helpers/data/embedding/small_static_embedding/' 'glove.6B.50d_test.txt', only_use_pretrain_word=True, lower=True) check_vector_equal(['The', 'Of', 'a'], vocab, embed, embed_dict, lower=True) check_word_unk(['notinfile'], vocab, embed) # 测试min_freq vocab = Vocabulary().add_word_lst(['The', 'a', 'notinfile1', 'A', 'notinfile2', 'notinfile2']) vocab.add_word('Of', no_create_entry=True) embed = StaticEmbedding(vocab, model_dir_or_name=tests_folder+'/helpers/data/embedding/small_static_embedding/' 'glove.6B.50d_test.txt', only_use_pretrain_word=True, lower=True, min_freq=2, only_train_min_freq=True) check_vector_equal(['Of', 'a'], vocab, embed, embed_dict, lower=True) check_word_unk(['notinfile1', 'The', 'notinfile2'], vocab, embed) def test_sequential_index(self): # 当不存在no_create_entry时，words_to_words应该是顺序的 vocab = Vocabulary().add_word_lst(['The', 'a', 'notinfile1', 'A', 'notinfile2', 'notinfile2']) embed = StaticEmbedding(vocab, model_dir_or_name=tests_folder+'/helpers/data/embedding/small_static_embedding/' 'glove.6B.50d_test.txt') for index,i in enumerate(embed.words_to_words): assert index==i embed_dict = read_static_embed(tests_folder+'/helpers/data/embedding/small_static_embedding/' 'glove.6B.50d_test.txt') for word, index in vocab: if word in embed_dict: index = vocab.to_index(word) v1 = embed(torch.LongTensor([index])).tolist()[0] v2 = embed_dict[word] for v1i, v2i in zip(v1, v2): assert np.allclose(v1i, v2i) def test_save_load_static_embed(self): static_test_folder = 'static_save_test' try: # 测试包含no_create_entry os.makedirs(static_test_folder, exist_ok=True) vocab = Vocabulary().add_word_lst(['The', 'a', 'notinfile1', 'A']) vocab.add_word_lst(['notinfile2', 'notinfile2'], no_create_entry=True) embed = StaticEmbedding(vocab, model_dir_or_name=tests_folder+'/helpers/data/embedding/small_static_embedding/' 'glove.6B.50d_test.txt') embed.save(static_test_folder) load_embed = StaticEmbedding.load(static_test_folder) words = torch.randint(len(vocab), size=(2, 20)) assert (embed(words) - load_embed(words)).sum() == 0 # 测试不包含no_create_entry vocab = Vocabulary().add_word_lst(['The', 'a', 'notinfile1', 'A']) embed = StaticEmbedding(vocab, model_dir_or_name=tests_folder+'/helpers/data/embedding/small_static_embedding/' 'glove.6B.50d_test.txt') embed.save(static_test_folder) load_embed = StaticEmbedding.load(static_test_folder) words = torch.randint(len(vocab), size=(2, 20)) assert (embed(words) - load_embed(words)).sum() == 0 # 测试lower, min_freq vocab = Vocabulary().add_word_lst(['The', 'the', 'the', 'A', 'a', 'B']) embed = StaticEmbedding(vocab, model_dir_or_name=tests_folder+'/helpers/data/embedding/small_static_embedding/' 'glove.6B.50d_test.txt', min_freq=2, lower=True) embed.save(static_test_folder) load_embed = StaticEmbedding.load(static_test_folder) words = torch.randint(len(vocab), size=(2, 20)) assert (embed(words) - load_embed(words)).sum() == 0 # 测试random的embedding vocab = Vocabulary().add_word_lst(['The', 'the', 'the', 'A', 'a', 'B']) vocab = vocab.add_word_lst(['b'], no_create_entry=True) embed = StaticEmbedding(vocab, model_dir_or_name=None, embedding_dim=4, min_freq=2, lower=True, normalize=True) embed.weight.data += 0.2 # 使得它不是normalize embed.save(static_test_folder) load_embed = StaticEmbedding.load(static_test_folder) words = torch.randint(len(vocab), size=(2, 20)) assert (embed(words) - load_embed(words)).sum()==0 finally: if os.path.isdir(static_test_folder): import shutil shutil.rmtree(static_test_folder) def read_static_embed(fp): """ :param str fp: embedding的路径 :return: {}, key是word, value是vector """ embed = {} with open(fp, 'r') as f: for line in f: line = line.strip() if line: parts = line.split() vector = list(map(float, parts[1:])) word = parts[0] embed[word] = vector return embed @pytest.mark.torch class TestRandomSameEntry: def test_same_vector(self): vocab = Vocabulary().add_word_lst(["The", "the", "THE", 'a', "A"]) embed = StaticEmbedding(vocab, model_dir_or_name=None, embedding_dim=5, lower=True) words = torch.LongTensor([[vocab.to_index(word) for word in ["The", "the", "THE", 'a', 'A']]]) words = embed(words) embed_0 = words[0, 0] for i in range(1, 3): assert torch.sum(embed_0==words[0, i]).eq(len(embed_0)) embed_0 = words[0, 3] for i in range(3, 5): assert torch.sum(embed_0 == words[0, i]).eq(len(embed_0)) def test_dropout_close(self): vocab = Vocabulary().add_word_lst(["The", "the", "THE", 'a', "A"]) embed = StaticEmbedding(vocab, model_dir_or_name=None, embedding_dim=5, lower=True, dropout=0.5, word_dropout=0.9) words = torch.LongTensor([[vocab.to_index(word) for word in ["The", "the", "THE", 'a', 'A']]]) embed.eval() words = embed(words) embed_0 = words[0, 0] for i in range(1, 3): assert torch.sum(embed_0 == words[0, i]).eq(len(embed_0)) embed_0 = words[0, 3] for i in range(3, 5): assert torch.sum(embed_0 == words[0, i]).eq(len(embed_0))