|
|
import pandas as pd |
|
|
from transformers import pipeline |
|
|
from transformers.pipelines.base import Pipeline |
|
|
from scipy.sparse import csr_matrix |
|
|
from typing import Mapping, List, Tuple, Any |
|
|
from bertopic.representation._base import BaseRepresentation |
|
|
|
|
|
|
|
|
class ZeroShotClassification(BaseRepresentation): |
|
|
""" Zero-shot Classification on topic keywords with candidate labels |
|
|
|
|
|
Arguments: |
|
|
candidate_topics: A list of labels to assign to the topics if they |
|
|
exceed `min_prob` |
|
|
model: A transformers pipeline that should be initialized as |
|
|
"zero-shot-classification". For example, |
|
|
`pipeline("zero-shot-classification", model="facebook/bart-large-mnli")` |
|
|
pipeline_kwargs: Kwargs that you can pass to the transformers.pipeline |
|
|
when it is called. NOTE: Use `{"multi_label": True}` |
|
|
to extract multiple labels for each topic. |
|
|
min_prob: The minimum probability to assign a candidate label to a topic |
|
|
|
|
|
Usage: |
|
|
|
|
|
```python |
|
|
from bertopic.representation import ZeroShotClassification |
|
|
from bertopic import BERTopic |
|
|
|
|
|
# Create your representation model |
|
|
candidate_topics = ["space and nasa", "bicycles", "sports"] |
|
|
representation_model = ZeroShotClassification(candidate_topics, model="facebook/bart-large-mnli") |
|
|
|
|
|
# Use the representation model in BERTopic on top of the default pipeline |
|
|
topic_model = BERTopic(representation_model=representation_model) |
|
|
``` |
|
|
""" |
|
|
def __init__(self, |
|
|
candidate_topics: List[str], |
|
|
model: str = "facebook/bart-large-mnli", |
|
|
pipeline_kwargs: Mapping[str, Any] = {}, |
|
|
min_prob: float = 0.8 |
|
|
): |
|
|
self.candidate_topics = candidate_topics |
|
|
if isinstance(model, str): |
|
|
self.model = pipeline("zero-shot-classification", model=model) |
|
|
elif isinstance(model, Pipeline): |
|
|
self.model = model |
|
|
else: |
|
|
raise ValueError("Make sure that the HF model that you" |
|
|
"pass is either a string referring to a" |
|
|
"HF model or a `transformers.pipeline` object.") |
|
|
self.pipeline_kwargs = pipeline_kwargs |
|
|
self.min_prob = min_prob |
|
|
|
|
|
def extract_topics(self, |
|
|
topic_model, |
|
|
documents: pd.DataFrame, |
|
|
c_tf_idf: csr_matrix, |
|
|
topics: Mapping[str, List[Tuple[str, float]]] |
|
|
) -> Mapping[str, List[Tuple[str, float]]]: |
|
|
""" Extract topics |
|
|
|
|
|
Arguments: |
|
|
topic_model: Not used |
|
|
documents: Not used |
|
|
c_tf_idf: Not used |
|
|
topics: The candidate topics as calculated with c-TF-IDF |
|
|
|
|
|
Returns: |
|
|
updated_topics: Updated topic representations |
|
|
""" |
|
|
|
|
|
topic_descriptions = [" ".join(list(zip(*topics[topic]))[0]) for topic in topics.keys()] |
|
|
classifications = self.model(topic_descriptions, self.candidate_topics, **self.pipeline_kwargs) |
|
|
|
|
|
|
|
|
updated_topics = {} |
|
|
for topic, classification in zip(topics.keys(), classifications): |
|
|
topic_description = topics[topic] |
|
|
|
|
|
|
|
|
if self.pipeline_kwargs.get("multi_label"): |
|
|
topic_description = [] |
|
|
for label, score in zip(classification["labels"], classification["scores"]): |
|
|
if score > self.min_prob: |
|
|
topic_description.append((label, score)) |
|
|
|
|
|
|
|
|
elif classification["scores"][0] > self.min_prob: |
|
|
topic_description = [(classification["labels"][0], classification["scores"][0])] |
|
|
|
|
|
|
|
|
if len(topic_description) == 0: |
|
|
topic_description = topics[topic] |
|
|
elif len(topic_description) < 10: |
|
|
topic_description += [("", 0) for _ in range(10-len(topic_description))] |
|
|
updated_topics[topic] = topic_description |
|
|
|
|
|
return updated_topics |
|
|
|
