Update README.md

README.md

@@ -73,4 +73,97 @@ With the help of a well-known AI, a data set was compiled to verify the veracity
 
 It can be seen from the above that the model is able to distinguish between basic emotions, but has problems with specific ones (grief (горе), pride (гордость)), and also defines a neutral emotion more often than expected (most likely this happened due to inaccuracies in the dataset, since in ordinary speech you can "use" several emojis at once and most often people use the neutral one)
 
-The model will be updated as improvements are made.
+The model will be updated as improvements are made.
+
+# Usage example
+```
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+tokenizer = AutoTokenizer.from_pretrained('AniMAntZeZo/RuBert-tiny2-EmotionsDetected')
+model = AutoModelForSequenceClassification.from_pretrained('AniMAntZeZo/RuBert-tiny2-EmotionsDetected')
+model.to("cuda" if torch.cuda.is_available() else "cpu",
+
+
+emotion_columns = [
+    "admiration", "amusement", "anger", "annoyance", "approval", "caring", "confusion", "curiosity", "desire",
+    "disappointment", "disapproval", "disgust", "embarrassment", "excitement", "fear", "gratitude", "grief", "joy",
+    "love", "nervousness", "optimism", "pride", "realization", "relief", "remorse", "sadness", "surprise", "neutral"
+]
+
+def predict_emotions(
+    text, 
+    model, 
+    tokenizer, 
+    emotion_columns, 
+    device="cuda" if torch.cuda.is_available() else "cpu", 
+    threshold=0.1
+):
+
+    emotion_translations = {
+        "admiration": "восхищение",
+        "amusement": "веселье",
+        "anger": "злость",
+        "annoyance": "раздражение",
+        "approval": "одобрение",
+        "caring": "забота",
+        "confusion": "непонимание",
+        "curiosity": "любопытство",
+        "desire": "желание",
+        "disappointment": "разочарование",
+        "disapproval": "неодобрение",
+        "disgust": "отвращение",
+        "embarrassment": "смущение",
+        "excitement": "возбуждение",
+        "fear": "страх",
+        "gratitude": "признательность",
+        "grief": "горе",
+        "joy": "радость",
+        "love": "любовь",
+        "nervousness": "нервозность",
+        "optimism": "оптимизм",
+        "pride": "гордость",
+        "realization": "осознание",
+        "relief": "облегчение",
+        "remorse": "раскаяние",
+        "sadness": "грусть",
+        "surprise": "удивление",
+        "neutral": "нейтральность",
+    }
+
+    model.to(device)
+    model.eval()
+    inputs = tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=128).to(device)
+    with torch.no_grad():
+        logits = model(**inputs).logits
+    probabilities = torch.sigmoid(logits).squeeze().cpu().numpy()
+
+    predictions = {
+        f"{emotion} ({emotion_translations[emotion]})": prob
+        for emotion, prob in zip(emotion_columns, probabilities) if prob > threshold
+    }
+
+    sorted_predictions = dict(sorted(predictions.items(), key=lambda item: item[1], reverse=True))
+    
+    return sorted_predictions
+```
+[INPUT]
+```
+example_text = "Как же я рад!"
+predictions = predict_emotions(example_text, model, tokenizer, emotion_columns)
+print("Emotions:", predictions)
+```
+[OUTPUT]
+```
+Emotions: {'joy (радость)': 0.6736836, 'excitement (возбуждение)': 0.25723574}
+```
+[INPUT]
+```
+example_text = "Я обиделся!"
+predictions = predict_emotions(example_text, model, tokenizer, emotion_columns)
+print("Emotions:", predictions)
+```
+[OUTPUT]
+```
+Emotions: {'sadness (грусть)': 0.3111033, 'disappointment (разочарование)': 0.2943853, 'annoyance (раздражение)': 0.19748639, 'anger (злость)': 0.16338393}
+```