Upload tool

pair_classification_tool.py

@@ -0,0 +1,23 @@
+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 = (
+        "This is a tool that classifies if two texts in English are similar or not using the labels 'equivalent' and "
+        "'not_equivalent'. It takes two inputs named `text` and `second_text` which should be in English and returns "
+        "the predicted label."
+    )
+
+    def encode(self, text, second_text):
+        return self.pre_processor(text, second_text, return_tensors="pt")
+
+    def decode(self, outputs):
+        logits = outputs.logits
+        label_id = torch.argmax(logits[0]).item()
+        return self.model.config.id2label[label_id]

tool_config.json

@@ -0,0 +1,7 @@
+{
+  "text_pair_classification": {
+    "description": "This is a tool that classifies if two texts in English are similar or not using the labels 'equivalent' and 'not_equivalent'. It takes two inputs named `text` and `second_text` which should be in English and returns the predicted label.",
+    "name": "pair_classifier",
+    "tool_class": "pair_classification_tool.TextPairClassificationTool"
+  }
+}