Spaces:
Running on Zero
Running on Zero
| import os | |
| import pandas as pd | |
| from datasets import load_dataset | |
| import sys | |
| from pathlib import Path | |
| # Add project root to path | |
| sys.path.append(str(Path(__file__).resolve().parent.parent.parent)) | |
| from src.config import TEXT_DATASET_PATH, DATA_DIR | |
| def download_goemotions(): | |
| print("[Dataset Expansion] Downloading GoEmotions dataset (43,410 text samples)...") | |
| try: | |
| # Load the dataset from HuggingFace using the full namespace to prevent URI errors | |
| dataset = load_dataset("google-research-datasets/go_emotions", "simplified") | |
| train_data = dataset["train"].to_pandas() | |
| print("[Dataset Expansion] GoEmotions downloaded successfully. Processing...") | |
| # GoEmotions has 27 labels. We will map them to your system's Stress Categories | |
| # 0: admiration, 1: amusement, 2: anger, 3: annoyance, 4: approval, 5: caring, | |
| # 6: confusion, 7: curiosity, 8: desire, 9: disappointment, 10: disapproval, | |
| # 11: disgust, 12: embarrassment, 13: excitement, 14: fear, 15: gratitude, | |
| # 16: grief, 17: joy, 18: love, 19: nervousness, 20: optimism, 21: pride, | |
| # 22: realization, 23: relief, 24: remorse, 25: sadness, 26: surprise, 27: neutral | |
| def map_emotion_to_stress(label_list): | |
| if not len(label_list): | |
| return "Normal" | |
| primary_label = label_list[0] | |
| # Stress: anger, annoyance, disapproval, disgust | |
| if primary_label in [2, 3, 10, 11]: | |
| return "Stress" | |
| # Depression: disappointment, grief, remorse, sadness | |
| elif primary_label in [9, 16, 24, 25]: | |
| return "Depression" | |
| # Anxiety: fear, nervousness | |
| elif primary_label in [14, 19]: | |
| return "Anxiety" | |
| # Emotional Distress: confusion, curiosity, desire, embarrassment, excitement, realization, surprise | |
| elif primary_label in [6, 7, 8, 12, 13, 22, 26]: | |
| return "Emotional Distress" | |
| # Normal: admiration, amusement, approval, caring, gratitude, joy, love, optimism, pride, relief, neutral | |
| else: | |
| return "Normal" | |
| train_data['category'] = train_data['labels'].apply(map_emotion_to_stress) | |
| # Create final dataframe | |
| expanded_df = pd.DataFrame({ | |
| "text": train_data["text"], | |
| "category": train_data["category"] | |
| }) | |
| # Append to existing dataset or save as new | |
| if os.path.exists(TEXT_DATASET_PATH): | |
| existing_df = pd.read_csv(TEXT_DATASET_PATH) | |
| # Make sure we only append if columns match | |
| if "text" in existing_df.columns and "category" in existing_df.columns: | |
| final_df = pd.concat([existing_df, expanded_df], ignore_index=True) | |
| final_df = final_df.drop_duplicates(subset=["text"]) | |
| final_df.to_csv(TEXT_DATASET_PATH, index=False) | |
| print(f"[Dataset Expansion] SUCCESS! Expanded Text Dataset to {len(final_df)} rows. Saved to {TEXT_DATASET_PATH}") | |
| else: | |
| out_path = os.path.join(DATA_DIR, "massive_text_dataset.csv") | |
| expanded_df.to_csv(out_path, index=False) | |
| print(f"[Dataset Expansion] Saved new massive dataset to {out_path}") | |
| else: | |
| expanded_df.to_csv(TEXT_DATASET_PATH, index=False) | |
| print(f"[Dataset Expansion] Created new dataset at {TEXT_DATASET_PATH}") | |
| except Exception as e: | |
| print(f"[Dataset Expansion] Error downloading GoEmotions: {e}") | |
| def audio_dataset_instructions(): | |
| print("\n" + "="*80) | |
| print("AUDIO DATASET EXPANSION (CMU-MOSEI / DAIC-WOZ)") | |
| print("="*80) | |
| print("Due to strict academic Non-Disclosure Agreements (NDAs) and massive file sizes (60GB+),") | |
| print("you must manually request access to DAIC-WOZ and D-Vlog from their university creators.") | |
| print("\nTo use the massive open-source CMU-MOSEI dataset, use the CMU Multimodal SDK in Colab:") | |
| print("1. Run: !pip install mmsdk") | |
| print("2. In a Colab cell, use the following code to download it to your Drive:") | |
| print(" from mmsdk import mmdatasdk") | |
| print(" cmumosei_highlevel = mmdatasdk.mmdataset(mmdatasdk.cmu_mosei.highlevel, '/content/drive/MyDrive/NeuroSense_AI/data/MOSEI/')") | |
| print("="*80 + "\n") | |
| if __name__ == "__main__": | |
| download_goemotions() | |
| audio_dataset_instructions() | |