Add BERTopic.

LLMTopicDetection_BERTopic/tests/test_sub_models/test_embeddings.py

@@ -0,0 +1,65 @@
+import copy
+import pytest
+import numpy as np
+from bertopic import BERTopic
+from sklearn.metrics.pairwise import cosine_similarity
+
+
+@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_extract_embeddings(model, request):
+    topic_model = copy.deepcopy(request.getfixturevalue(model))
+    single_embedding = topic_model._extract_embeddings("a document")
+    multiple_embeddings = topic_model._extract_embeddings(["something different", "another document"])
+    sim_matrix = cosine_similarity(single_embedding, multiple_embeddings)[0]
+
+    assert single_embedding.shape[0] == 1
+    assert single_embedding.shape[1] == 384
+    assert np.min(single_embedding) > -5
+    assert np.max(single_embedding) < 5
+
+    assert multiple_embeddings.shape[0] == 2
+    assert multiple_embeddings.shape[1] == 384
+    assert np.min(multiple_embeddings) > -5
+    assert np.max(multiple_embeddings) < 5
+
+    assert sim_matrix[0] < 0.5
+    assert sim_matrix[1] > 0.5
+
+
+@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_extract_embeddings_compare(model, embedding_model, request):
+    docs = ["some document"]
+    topic_model = copy.deepcopy(request.getfixturevalue(model))
+    bertopic_embeddings = topic_model._extract_embeddings(docs)
+
+    assert isinstance(bertopic_embeddings, np.ndarray)
+    assert bertopic_embeddings.shape == (1, 384)
+
+    sentence_embeddings = embedding_model.encode(docs, show_progress_bar=False)
+    assert np.array_equal(bertopic_embeddings, sentence_embeddings)
+
+
+def test_extract_incorrect_embeddings():
+    with pytest.raises(ValueError):
+        model = BERTopic(language="Unknown language")
+        model.fit(["some document"])