File size: 3,622 Bytes
3733c61
 
 
 
 
e6f08d5
3733c61
 
e6f08d5
3733c61
 
 
965d118
 
 
 
 
e6f08d5
3733c61
 
e6f08d5
3733c61
 
 
e6f08d5
3733c61
e6f08d5
3733c61
965d118
 
e6f08d5
965d118
 
3733c61
e6f08d5
965d118
 
3733c61
e6f08d5
3733c61
e6f08d5
3733c61
 
 
 
 
 
 
 
e6f08d5
3733c61
 
 
 
 
 
 
e6f08d5
3733c61
 
 
e6f08d5
3733c61
 
e6f08d5
3733c61
 
 
e6f08d5
3733c61
 
e6f08d5
3733c61
 
e6f08d5
3733c61
e6f08d5
 
3733c61
 
 
 
 
 
e6f08d5
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
import os
import torch
import torch.nn.functional as F
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification

SENTIMENT_LABELS = ['Negative', 'Neutral', 'Positive']
EMOTION_LABELS = ['Joy', 'Anger', 'Fear', 'Sadness']

class SentimentEmotionPredictor:
    def __init__(self):
        base_dir = os.path.dirname(os.path.abspath(__file__))
        local_sentiment = os.path.join(base_dir, "models", "sentiment_model")
        local_emotion = os.path.join(base_dir, "models", "emotion_model")

        self.sentiment_path = local_sentiment if os.path.exists(local_sentiment) else "usman-ai-dev/urdu-sentiment-xlmr"
        self.emotion_path = local_emotion if os.path.exists(local_emotion) else "usman-ai-dev/urdu-emotion-xlmr"

        self.sentiment_map = {0: "Negative", 1: "Neutral", 2: "Positive"}
        self.emotion_map = {0: "Joy", 1: "Anger", 2: "Fear", 3: "Sadness"}

        self.tokenizer = None
        self.sentiment_model = None
        self.emotion_model = None

        self.load_models()

    def load_models(self):
        print(f"Loading Tokenizer from '{self.sentiment_path}'...")
        self.tokenizer = AutoTokenizer.from_pretrained(self.sentiment_path)

        print(f"Loading Sentiment Model from '{self.sentiment_path}'...")
        self.sentiment_model = AutoModelForSequenceClassification.from_pretrained(self.sentiment_path)
        self.sentiment_model.eval()

        print(f"Loading Emotion Model from '{self.emotion_path}'...")
        self.emotion_model = AutoModelForSequenceClassification.from_pretrained(self.emotion_path)
        self.emotion_model.eval()

        print("Models loaded into memory successfully!")

    def get_word_attention(self, text):
        inputs = self.tokenizer(text, return_tensors="pt", max_length=128, truncation=True)
        with torch.no_grad():
            outputs = self.sentiment_model(**inputs, output_attentions=True)
        
        # Last layer attention averaged across heads for the [CLS] token
        attn = outputs.attentions[-1].mean(dim=1).squeeze(0)[0, :]
        tokens = self.tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])

        attention_list = []
        for tok, score in zip(tokens, attn):
            if tok not in ['<s>', '</s>', '<pad>']:
                clean_tok = tok.replace(' ', '') if tok.startswith(' ') else tok
                if clean_tok:
                    attention_list.append({'word': clean_tok, 'score': round(float(score), 4)})
        return attention_list

    def predict(self, text):
        if not text or not str(text).strip():
            return {"error": "Empty text provided."}

        text_str = str(text).strip()
        inputs = self.tokenizer(text_str, return_tensors="pt", truncation=True, max_length=128)

        with torch.no_grad():
            s_outputs = self.sentiment_model(**inputs, output_attentions=True)
            e_outputs = self.emotion_model(**inputs)

        s_probs = F.softmax(s_outputs.logits, dim=-1)[0]
        e_probs = F.softmax(e_outputs.logits, dim=-1)[0]

        s_idx = int(torch.argmax(s_probs).item())
        e_idx = int(torch.argmax(e_probs).item())

        attention_scores = self.get_word_attention(text_str)

        return {
            "text": text_str,
            "sentiment": self.sentiment_map[s_idx],
            "sentiment_scores": {l: round(float(p), 4) for l, p in zip(SENTIMENT_LABELS, s_probs)},
            "emotion": self.emotion_map[e_idx],
            "emotion_scores": {l: round(float(p), 4) for l, p in zip(EMOTION_LABELS, e_probs)},
            "attention": attention_scores
        }