def test_embed_one_sentence(): X = ["This is a sentence"] for embedding in EMBEDDING_TYPES: vec = bert_vectorizer.BertVectorizer(sentence_embedding=embedding) X_embed = vec.fit_transform(X) assert (X_embed.shape == (1, 768))
def test_embed_scibert(): X = ["This is a sentence" * 100] for embedding in EMBEDDING_TYPES: vec = bert_vectorizer.BertVectorizer(pretrained='scibert', sentence_embedding=embedding) X_embed = vec.fit_transform(X) assert (X_embed.shape == (1, 768))
def test_save_and_load(tmpdir): tmpfile = tmpdir.join('test.npy') X = ["This is a sentence" * 100] for pretrained in ['bert', 'scibert']: for embedding in EMBEDDING_TYPES: vec = bert_vectorizer.BertVectorizer(pretrained=pretrained, sentence_embedding=embedding) X_embed = vec.fit_transform(X) vec.save_transformed(str(tmpfile), X_embed) X_loaded = vec.load_transformed(str(tmpfile)) assert (X_loaded != X_embed).sum() == 0