| 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 |
|
|
