| import torch |
| from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer |
| from transformers.tools import PipelineTool |
|
|
|
|
| class TextPairClassificationTool(PipelineTool): |
| default_checkpoint = "sgugger/bert-finetuned-mrpc" |
| pre_processor_class = AutoTokenizer |
| model_class = AutoModelForSequenceClassification |
|
|
| description = ( |
| "classifies if two texts in English are similar or not using the labels {labels}. It takes two inputs named " |
| "`text` and `second_text` which should be in English and returns a dictionary with two keys named 'label' " |
| "(the predicted label ) and 'score' (the probability associated to it)." |
| ) |
|
|
| def post_init(self): |
| if isinstance(self.model, str): |
| config = AutoConfig.from_pretrained(self.model) |
| else: |
| config = self.model.config |
|
|
| labels = list(config.label2id.keys()) |
|
|
| if len(labels) > 1: |
| labels = [f"'{label}'" for label in labels] |
| labels_string = ", ".join(labels[:-1]) |
| labels_string += f", and {labels[-1]}" |
| else: |
| raise ValueError("Not enough labels.") |
|
|
| self.description = self.description.replace("{labels}", labels_string) |
|
|
| def encode(self, text, second_text): |
| return self.pre_processor(text, second_text, return_tensors="pt") |
|
|
| def decode(self, outputs): |
| logits = outputs.logits |
| scores = torch.nn.functional.softmax(logits, dim=-1) |
| label_id = torch.argmax(logits[0]).item() |
| label = self.model.config.id2label[label_id] |
| return {"label": label, "score": scores[0][label_id].item()} |
|
|
