Update README.md

README.md

@@ -26,6 +26,7 @@ dataset for multilabel classification. Model can be used to extract all emotions
 # Usage
 Using model with Huggingface Transformers:
 ```python
+import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 tokenizer = AutoTokenizer.from_pretrained("fyaronskiy/ruRoberta-large-ru-go-emotions")
 model = AutoModelForSequenceClassification.from_pretrained("fyaronskiy/ruRoberta-large-ru-go-emotions")
@@ -42,13 +43,12 @@ def predict_emotions(text):
   inputs = tokenizer(text, truncation=True, add_special_tokens=True, max_length=128, return_tensors='pt')
   with torch.no_grad():
       logits = model(**inputs).logits
-  probas = torch.sigmoid(logits).squeeze(dim=0)
-  probas = probas.cpu().numpy()
-  
-  class_binary_labels = (probas > np.array(best_thresholds)).astype(int)
+  probas = torch.sigmoid(logits).squeeze(dim=0)  
+  class_binary_labels = (probas > torch.tensor(best_thresholds)).int()
   return [ID2LABEL[label_id] for label_id, value in enumerate(class_binary_labels) if value == 1]
 
 print(predict_emotions('У вас отличный сервис и лучший кофе в городе, обожаю вашу кофейню!'))
+
 #['admiration', 'love']
 ```