Skip to content

Commit 6a2834d

Browse files
committed
Better backwards compatibility
1 parent 7283206 commit 6a2834d

5 files changed

Lines changed: 166 additions & 13 deletions

File tree

‎bertopic/_bertopic.py‎

Lines changed: 32 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -327,21 +327,18 @@ def probabilities_(self) -> list[int]:
327327
@property
328328
def topic_representations_(self) -> dict[int, list[tuple[str, float]]]:
329329
"""For backwards compatibility."""
330-
representations = {}
331-
for topic in self._topics:
332-
representations[topic.id] = topic.representations["Main"].data
333-
return representations if representations else None
330+
return {topic.id: topic.representations["Main"].word_scores for topic in self._topics} or None
334331

335332
@property
336333
def topic_aspects_(self) -> dict[int, dict[str, list[tuple[str, float]]]]:
337334
"""For backwards compatibility."""
338335
first_topic = next(iter(self._topics), None)
339-
aspect_names = list(first_topic.representations) if first_topic else []
340-
aspects = {
341-
aspect_name: {topic.id: topic.representations[aspect_name].data for topic in self._topics}
342-
for aspect_name in aspect_names
336+
aspects = first_topic.representations if first_topic else []
337+
return {
338+
name: {topic.id: topic.representations[name].word_scores for topic in self._topics}
339+
for name in aspects
340+
if name != "Main"
343341
}
344-
return aspects
345342

346343
@property
347344
def topic_sizes_(self) -> dict[int, int]:
@@ -364,9 +361,9 @@ def topic_labels_(self) -> dict[int, str]:
364361
return self._topics.labels
365362

366363
@property
367-
def custom_labels_(self) -> list[str]:
368-
"""For backwards compatibility."""
369-
return [topic.label for topic in self._topics]
364+
def custom_labels_(self) -> list[str] | None:
365+
"""For backwards compatibility. None until `set_topic_labels` is called."""
366+
return self._topics.custom_labels
370367

371368
@property
372369
def representative_docs_(self) -> dict[int, str]:
@@ -1129,6 +1126,29 @@ def get_topic(self, topic: int, full: bool = False) -> dict[str, tuple[str, floa
11291126
else:
11301127
return False
11311128

1129+
def get_representation(self, topic: int, aspect: str = "Main") -> TopicRepresentation:
1130+
"""Return a topic's representation with its structure intact.
1131+
1132+
`get_topic` flattens every representation to words and scores, which loses the shape
1133+
of anything richer than keywords, such as the dictionary an LLM returns. This returns
1134+
the representation itself.
1135+
1136+
Arguments:
1137+
topic: A specific topic for which you want its representation
1138+
aspect: Which representation to return, defaulting to the main one
1139+
1140+
Returns:
1141+
The `TopicRepresentation` for that topic and aspect
1142+
1143+
Examples:
1144+
```python
1145+
representation = topic_model.get_representation(12)
1146+
representation.words
1147+
```
1148+
"""
1149+
check_is_fitted(self)
1150+
return self._topics[topic].representations[aspect]
1151+
11321152
def get_topic_info(self, topic: int | None = None) -> pl.DataFrame:
11331153
"""Get information about each topic including its ID, frequency, and name.
11341154

‎bertopic/_topics.py‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,20 @@ def words(self) -> list[str]:
2828
"""
2929
return []
3030

31+
@property
32+
def scores(self) -> list[float]:
33+
"""One score per word, weighting every word equally.
34+
35+
Only `Keywords` ranks the words within a topic. Everything else names the topic
36+
instead of weighting what is in it, so no word is more representative than another.
37+
"""
38+
return [1.0] * len(self.words)
39+
40+
@property
41+
def word_scores(self) -> list[tuple[str, float]]:
42+
"""Words paired with their scores, which is how `get_topic` returns a topic."""
43+
return list(zip(self.words, self.scores))
44+
3145
def __str__(self) -> str:
3246
"""String representation of the topic representation."""
3347
return str(self.data)
@@ -464,6 +478,17 @@ def labels(self) -> dict[int, str]:
464478
labels = {topic.id: topic.label if topic.label is not None else "" for topic in self.topics.values()}
465479
return dict(sorted(labels.items()))
466480

481+
@property
482+
def custom_labels(self) -> list[str] | None:
483+
"""The labels a user set, or None when every label is still derived.
484+
485+
`Topic._label` is only filled in by `set_topic_labels`, which is what separates
486+
a label the user chose from one generated out of the topic's keywords.
487+
"""
488+
if all(topic._label is None for topic in self):
489+
return None
490+
return [topic.label for topic in self]
491+
467492
@property
468493
def c_tf_idf(self) -> csr_matrix:
469494
"""Get c-TF-IDF matrix for all topics in ascending topic ID order

‎tests/test_representation/test_get.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
import pytest
33
import polars as pl
44

5+
from bertopic._topics import Label, StructuredJSON
6+
57

68
@pytest.mark.parametrize(
79
"model",
@@ -123,3 +125,60 @@ def test_get_topic_info(model, request):
123125
assert len(topic_model.get_topic_info(topic)) == 1
124126

125127
assert len(topic_model.get_topic_info(200)) == 0
128+
129+
130+
@pytest.mark.parametrize(
131+
"model",
132+
[
133+
("kmeans_pca_topic_model"),
134+
("base_topic_model"),
135+
("custom_topic_model"),
136+
("merged_topic_model"),
137+
("reduced_topic_model"),
138+
("online_topic_model"),
139+
],
140+
)
141+
def test_get_topic_always_returns_words_with_scores(model, request):
142+
topic_model = copy.deepcopy(request.getfixturevalue(model))
143+
144+
for topic in set(topic_model.topics_):
145+
representation = topic_model.get_topic(topic)
146+
assert all(isinstance(word, str) and isinstance(score, float) for word, score in representation)
147+
148+
149+
def test_get_topic_flattens_a_label_representation(base_topic_model):
150+
"""The regression this unit exists for: a label used to come back as a bare string.
151+
152+
`get_topic` returned `'label_0'`, so `visualize_barchart` raised on `for word, _ in ...`
153+
and `_create_topic_vectors` would have embedded individual characters.
154+
"""
155+
topic_model = copy.deepcopy(base_topic_model)
156+
topic_ids = topic_model._topics.topic_ids()
157+
topic_model._topics.set_data(
158+
representations={"Main": {topic_id: Label(data=f"label_{topic_id}") for topic_id in topic_ids}}
159+
)
160+
first = topic_ids[0]
161+
162+
assert topic_model.get_topic(first) == [(f"label_{first}", 1.0)]
163+
assert topic_model.visualize_barchart() is not None
164+
165+
166+
def test_get_representation_keeps_the_structure_get_topic_flattens(base_topic_model):
167+
"""Flattening is lossy, so the representation itself stays reachable."""
168+
topic_model = copy.deepcopy(base_topic_model)
169+
first = topic_model._topics.topic_ids()[0]
170+
topic_model._topics.set_data(representations={"Main": {first: StructuredJSON(data={"topic": "cats"})}})
171+
172+
representation = topic_model.get_representation(first)
173+
174+
assert isinstance(representation, StructuredJSON)
175+
assert representation.data == {"topic": "cats"}
176+
assert topic_model.get_topic(first) == [("cats", 1.0)]
177+
178+
179+
def test_topic_aspects_excludes_the_main_representation(representation_topic_model):
180+
"""Aspects are the additional representations; Main is reached through `get_topic`."""
181+
aspects = representation_topic_model.topic_aspects_
182+
183+
assert "Main" not in aspects
184+
assert "MMR" in aspects

‎tests/test_representation/test_labels.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,9 @@ def test_generate_topic_labels(model, request):
4040
def test_set_labels(model, request):
4141
topic_model = copy.deepcopy(request.getfixturevalue(model))
4242

43+
# Nothing is custom until the user says so, which is what the plotting guards test for
44+
assert topic_model.custom_labels_ is None
45+
4346
labels = topic_model.generate_topic_labels()
4447
topic_model.set_topic_labels(labels)
4548
assert topic_model.custom_labels_ == labels

‎tests/test_topics.py‎

Lines changed: 47 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,17 @@
2424
from scipy.sparse import csr_matrix
2525

2626
from bertopic._corpus import Corpus
27-
from bertopic._topics import Images, Keywords, Topic, TopicHierarchy, TopicMapping, Topics, TopicType
27+
from bertopic._topics import (
28+
Images,
29+
Keywords,
30+
Label,
31+
StructuredJSON,
32+
Topic,
33+
TopicHierarchy,
34+
TopicMapping,
35+
Topics,
36+
TopicType,
37+
)
2838

2939

3040
def build_topics(counts: dict[int, int], probabilities: np.ndarray | None = None) -> Topics:
@@ -403,6 +413,42 @@ def test_mapping_rejects_an_incomplete_new_mapping():
403413
# --------------------------------------------------------------------------------------
404414

405415

416+
def test_keywords_carry_their_own_scores():
417+
"""Keywords are the one representation that ranks the words within a topic."""
418+
keywords = Keywords(data=[("cat", 0.8), ("dog", 0.4)])
419+
420+
assert keywords.words == ["cat", "dog"]
421+
assert keywords.scores == [0.8, 0.4]
422+
assert keywords.word_scores == [("cat", 0.8), ("dog", 0.4)]
423+
424+
425+
@pytest.mark.parametrize(
426+
"representation,expected_words",
427+
[
428+
(Label(data="animals"), ["animals"]),
429+
(StructuredJSON(data={"topic": "animals", "theme": "pets"}), ["animals", "pets"]),
430+
(Images(data="collage", captions=["a cat", "a dog"]), ["a cat", "a dog"]),
431+
],
432+
)
433+
def test_representations_without_scores_weight_every_word_equally(representation, expected_words):
434+
"""`words` reads any representation as text, and `scores` is what ranks those words.
435+
436+
Only keywords rank them. Everything else names the topic, so each word counts the same
437+
and `word_scores` is what `get_topic` hands back.
438+
"""
439+
assert representation.words == expected_words
440+
assert representation.scores == [1.0] * len(expected_words)
441+
assert representation.word_scores == [(word, 1.0) for word in expected_words]
442+
443+
444+
def test_a_representation_with_no_text_view_has_no_scores():
445+
"""Scores follow words, so a payload with no text view yields neither."""
446+
images = Images(data="collage")
447+
448+
assert images.words == []
449+
assert images.word_scores == []
450+
451+
406452
def test_images_have_no_words_until_they_are_captioned():
407453
"""A representation may hold a modality that has no text view yet."""
408454
images = Images(data="collage")

0 commit comments

Comments
 (0)