Update README.md

README.md

@@ -11,11 +11,15 @@ from transformers import AutoModelForSequenceClassification, BertTokenizer
 model_name = 'cointegrated/rubert-base-cased-dp-paraphrase-detection'
 model = AutoModelForSequenceClassification.from_pretrained(model_name).cuda()
 tokenizer = BertTokenizer.from_pretrained(model_name)
-text1 = 'Сегодня на улице хорошая погода'
-text2 = 'Сегодня на улице отвратительная погода'
-batch = tokenizer(text1, text2, return_tensors='pt').to(model.device)
-with torch.inference_mode():
-    proba = torch.softmax(model(**batch).logits, -1).cpu().numpy()
-print(proba)
-# [[0.44876656 0.5512334 ]]
+
+def compare_texts(text1, text2):
+    batch = tokenizer(text1, text2, return_tensors='pt').to(model.device)
+    with torch.inference_mode():
+        proba = torch.softmax(model(**batch).logits, -1).cpu().numpy()
+    return proba[0] # p(non-paraphrase), p(paraphrase)
+
+print(compare_texts('Сегодня на улице хорошая погода', 'Сегодня на улице отвратительная погода'))
+# [0.7056226 0.2943774]
+print(compare_texts('Сегодня на улице хорошая погода', 'Отличная погодка сегодня выдалась'))
+# [0.16524374 0.8347562 ]
 ```