Update README.md

README.md

@@ -17,24 +17,19 @@ This is a baseline RoBERTa-base model for the delicate text detection task.
 Here's a short usage example with the torch library in a binary classification task:
 
 ```python
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import torch
+from transformers import pipeline
 
-tokenizer = AutoTokenizer.from_pretrained("grammarly/detexd-roberta")
-model = AutoModelForSequenceClassification.from_pretrained("grammarly/detexd-roberta")
-model.eval()
+classifier = pipeline("text-classification", model="grammarly/detexd-roberta-base")
 
-def predict_binary_score(text: str, break_class_ix=3):
-    with torch.no_grad():
-        # get multiclass probability scores
-        logits = model(**tokenizer(text, return_tensors='pt'))[0]
-        probs = torch.nn.functional.softmax(logits, dim=-1)
+def predict_binary_score(text: str):
+    # get multiclass probability scores
+    scores = classifier(text, top_k=None)
 
-        # convert to a binary prediction by summing the probability scores
-        # for the higher-index classes, as defined by break_class_ix
-        bin_score = probs[..., break_class_ix:].sum(dim=-1)
-
-        return bin_score.item()
+    # convert to a single score by summing the probability scores
+    # for the higher-index classes
+    return sum(score['score']
+               for score in scores
+               if score['label'] in ('LABEL_3', 'LABEL_4', 'LABEL_5'))
 
 def predict_delicate(text: str, threshold=0.72496545):
     return predict_binary_score(text) > threshold