Spaces:
Sleeping
Sleeping
| """ | |
| 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"] | |
| } | |