File size: 1,387 Bytes
19b102a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
import copy
import pytest
import numpy as np
from umap import UMAP
from sklearn.decomposition import PCA
from bertopic import BERTopic
@pytest.mark.parametrize("dim_model", [UMAP, PCA])
@pytest.mark.parametrize("embeddings,shape,n_components", [(np.random.rand(100, 128), 100, 5),
(np.random.rand(10, 256), 10, 5),
(np.random.rand(50, 15), 50, 10)])
def test_reduce_dimensionality(dim_model, embeddings, shape, n_components):
model = BERTopic(umap_model=dim_model(n_components=n_components))
umap_embeddings = model._reduce_dimensionality(embeddings)
assert umap_embeddings.shape == (shape, n_components)
@pytest.mark.parametrize('model', [('kmeans_pca_topic_model'),
('base_topic_model'),
('custom_topic_model'),
('merged_topic_model'),
('reduced_topic_model'),
('online_topic_model')])
def test_custom_reduce_dimensionality(model, request):
embeddings = np.random.rand(500, 128)
topic_model = copy.deepcopy(request.getfixturevalue(model))
umap_embeddings = topic_model._reduce_dimensionality(embeddings)
assert umap_embeddings.shape[1] < embeddings.shape[1]
|