Update README.md

README.md

@@ -70,14 +70,42 @@ pt
 from transformers import XLMRobertaTokenizerFast, XLMRobertaForSequenceClassification
 
 # load tokenizer and model weights
-tokenizer = XLMRobertaTokenizerFast.from_pretrained('SkolkovoInstitute/xlmr_formality_classifier')
-model = XLMRobertaForSequenceClassification.from_pretrained('SkolkovoInstitute/xlmr_formality_classifier')
+tokenizer = XLMRobertaTokenizerFast.from_pretrained('s-nlp/xlmr_formality_classifier')
+model = XLMRobertaForSequenceClassification.from_pretrained('s-nlp/xlmr_formality_classifier')
+
+id2formality = {0: "formal", 1: "informal"}
+texts = [
+    "I like you. I love you",
+    "Hey, what's up?",
+    "Siema, co porabiasz?",
+    "I feel deep regret and sadness about the situation in international politics.",
+]
 
 # prepare the input
-batch = tokenizer.encode('ты супер', return_tensors='pt')
+encoding = tokenizer(
+    texts,
+    add_special_tokens=True,
+    return_token_type_ids=True,
+    truncation=True,
+    padding="max_length",
+    return_tensors="pt",
+)
 
 # inference
-model(batch)
+output = model(**encoding)
+
+formality_scores = [
+    {id2formality[idx]: score for idx, score in enumerate(text_scores.tolist())}
+    for text_scores in output.logits.softmax(dim=1)
+]
+formality_scores
+```
+
+```
+[{'formal': 0.993225634098053, 'informal': 0.006774314679205418},
+ {'formal': 0.8807966113090515, 'informal': 0.1192033663392067},
+ {'formal': 0.936184287071228, 'informal': 0.06381577253341675},
+ {'formal': 0.9986615180969238, 'informal': 0.0013385231141000986}]
 ```
 
 ## Citation