File size: 4,927 Bytes
feaf7eb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
"""
Configuration for different emotion detection models
Add new models here with their specific settings
"""

MODELS_CONFIG = {
    # SuperB Wav2Vec2 - Lightweight, 4 emotions
    "superb/wav2vec2-base-superb-er": {
        "task": "audio-classification",
        "emotions": ["Neutral", "Happy", "Sad", "Angry"],
        "label_mapping": {
            "neu": "Neutral",
            "neutral": "Neutral",
            "hap": "Happy",
            "happy": "Happy",
            "sad": "Sad",
            "sadness": "Sad",
            "ang": "Angry",
            "angry": "Angry",
            "anger": "Angry"
        },
        "sample_rate": 16000,
        "description": "Lightweight model with 4 basic emotions"
    },
    
    # SuperB HuBERT - Better accuracy, 4 emotions
    "superb/hubert-large-superb-er": {
        "task": "audio-classification",
        "emotions": ["Neutral", "Happy", "Sad", "Angry"],
        "label_mapping": {
            "neu": "Neutral",
            "neutral": "Neutral",
            "hap": "Happy",
            "happy": "Happy",
            "sad": "Sad",
            "sadness": "Sad",
            "ang": "Angry",
            "angry": "Angry",
            "anger": "Angry"
        },
        "sample_rate": 16000,
        "description": "HuBERT-based model with better accuracy"
    },
    
    # Ehcalabres Wav2Vec2 XLSR - 7 emotions
    "ehcalabres/wav2vec2-lg-xlsr-en-speech-emotion-recognition": {
        "task": "audio-classification",
        "emotions": ["Neutral", "Happy", "Sad", "Angry", "Fear", "Disgust", "Surprise"],
        "label_mapping": {
            "neu": "Neutral",
            "neutral": "Neutral",
            "hap": "Happy",
            "happy": "Happy",
            "happiness": "Happy",
            "sad": "Sad",
            "sadness": "Sad",
            "ang": "Angry",
            "angry": "Angry",
            "anger": "Angry",
            "fea": "Fear",
            "fear": "Fear",
            "dis": "Disgust",
            "disgust": "Disgust",
            "sur": "Surprise",
            "surprise": "Surprise"
        },
        "sample_rate": 16000,
        "description": "Multi-lingual model with 7 emotions"
    },
    
    # Harshit345 XLSR - Alternative model
    "harshit345/xlsr-wav2vec-speech-emotion-recognition": {
        "task": "automatic-speech-recognition",  # Different task type
        "emotions": ["Neutral", "Happy", "Sad", "Angry", "Fear", "Disgust", "Surprise"],
        "label_mapping": {
            "neutral": "Neutral",
            "calm": "Neutral",
            "happy": "Happy",
            "sad": "Sad",
            "angry": "Angry",
            "fearful": "Fear",
            "fear": "Fear",
            "disgust": "Disgust",
            "surprised": "Surprise",
            "surprise": "Surprise"
        },
        "sample_rate": 16000,
        "description": "XLSR-based emotion recognition",
        "special_handling": True  # Needs custom loading
    },
    
    # Amiriparian Wav2Vec2 - RAVDESS dataset
    "amiriparian/wav2vec2-base-ravdess": {
        "task": "audio-classification",
        "emotions": ["Neutral", "Happy", "Sad", "Angry", "Fear", "Disgust", "Surprise", "Calm"],
        "label_mapping": {
            "01": "Neutral",
            "02": "Calm",
            "03": "Happy",
            "04": "Sad",
            "05": "Angry",
            "06": "Fear",
            "07": "Disgust",
            "08": "Surprise",
            "neutral": "Neutral",
            "calm": "Calm",
            "happy": "Happy",
            "sad": "Sad",
            "angry": "Angry",
            "fearful": "Fear",
            "fear": "Fear",
            "disgust": "Disgust",
            "surprised": "Surprise",
            "surprise": "Surprise"
        },
        "sample_rate": 16000,
        "description": "Trained on RAVDESS dataset with 8 emotions"
    }
}

def get_model_config(model_name):
    """
    Get configuration for a specific model
    
    Args:
        model_name: Name of the model
    
    Returns:
        dict: Model configuration or default config
    """
    if model_name in MODELS_CONFIG:
        return MODELS_CONFIG[model_name]
    
    # Default configuration for unknown models
    return {
        "task": "audio-classification",
        "emotions": ["Neutral", "Happy", "Sad", "Angry"],
        "label_mapping": {},
        "sample_rate": 16000,
        "description": "Custom model",
        "special_handling": False
    }

def get_available_models():
    """Get list of all available configured models"""
    return list(MODELS_CONFIG.keys())

def get_model_info(model_name):
    """Get human-readable info about a model"""
    config = get_model_config(model_name)
    return {
        "name": model_name,
        "emotions": config["emotions"],
        "num_emotions": len(config["emotions"]),
        "description": config["description"],
        "sample_rate": config["sample_rate"]
    }