diff --git a/.dockerignore b/.dockerignore
new file mode 100644
index 0000000000000000000000000000000000000000..c836eec753309da41c79cb365b401de3c063df99
--- /dev/null
+++ b/.dockerignore
@@ -0,0 +1,32 @@
+# Large data directories
+data/
+venv/
+.venv/
+
+# Frontend (not needed for API server)
+app/
+node_modules/
+
+# Dev files
+notebooks/
+.git/
+.github/
+__pycache__/
+*.pyc
+
+# Archives
+*.zip
+
+# IDE
+.vscode/
+.idea/
+
+# Claude/gstack
+.claude/
+.gstack/
+.superpowers/
+
+# Docs (not needed at runtime)
+docs/
+*.md
+!requirements-deploy.txt
diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 0000000000000000000000000000000000000000..d899f6551a51cf19763c5955c7a06a2726f018e9
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1 @@
+*.wav filter=lfs diff=lfs merge=lfs -text
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 0000000000000000000000000000000000000000..b5ac03a56762ab789c9d75e9a660391b682536a3
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1,39 @@
+FROM python:3.12-slim
+
+# System deps for audio processing (librosa, soundfile)
+RUN apt-get update && apt-get install -y --no-install-recommends \
+ libsndfile1 \
+ ffmpeg \
+ && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+# Install torch + torchaudio CPU-only (saves ~2.3GB vs full CUDA build)
+RUN pip install --no-cache-dir \
+ torch torchaudio --index-url https://download.pytorch.org/whl/cpu
+
+# Force cache invalidation
+ARG CACHEBUST=6
+# Install remaining dependencies
+COPY requirements-deploy.txt .
+RUN pip install --no-cache-dir -r requirements-deploy.txt
+
+# Copy source code
+COPY src/ src/
+COPY config.yaml .
+COPY scripts/cache_models.py scripts/cache_models.py
+
+# Pre-download ML models at build time (avoids cold-start downloads)
+ARG HF_TOKEN
+ENV HF_TOKEN=${HF_TOKEN}
+RUN python scripts/cache_models.py
+
+# Create data directory for SQLite + uploads
+RUN mkdir -p data/samples
+
+# HF Spaces port
+ENV PORT=7860
+
+EXPOSE 7860
+
+CMD ["uvicorn", "src.stage4.main:app", "--host", "0.0.0.0", "--port", "7860"]
diff --git a/README.md b/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..73e262e8256e47a996061bdce4c9c7e163084a63
--- /dev/null
+++ b/README.md
@@ -0,0 +1,334 @@
+---
+title: UsTwo API
+emoji: ๐
+colorFrom: pink
+colorTo: purple
+sdk: docker
+app_port: 7860
+pinned: false
+---
+
+# UsTwo โ your characters react to your calls
+
+> **CMU 11-775 Large Scale Multimedia Analysis** โ Team Project
+> 3-person team: Juhyun (PO / Research) ยท Seungjae (ML) ยท Youngkyun (App)
+
+UsTwo takes a recorded phone call between two people (couple, friends, family) and runs **multimodal analysis (audio + text)** to understand how each speaker felt. It then produces three things inside a **React Native + FastAPI** mobile app: a **character reaction scene**, a **growing emotion garden**, and an **LLM-written recap card**.
+
+The goal isn't just emotion classification. It's to visualize *"what kind of moment did these two share on this call?"*
+
+
+ 
+ 
+ 
+
+
+---
+
+## At a glance
+
+```
+[call recording .wav/.m4a]
+ โ
+ โผ
+โโโโโโโโโโโโโโโโโโโโโ Stage 1 (Seungjae) โโโโโโโโโโโโโโโโโโโโโ
+โ pyannote 4.x (3.1) โ WhisperX large-v3-turbo INT8 โ ko/en โ
+โ speaker diarization ASR + forced alignment LID โ
+โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+ โ segments: [{speaker, start, end, text, lang}]
+ โผ
+โโโโโโโโโโโโโโโโโโโโโ Stage 2 (Seungjae) โโโโโโโโโโโโโโโโโโโโโ
+โ emotion2vec LoRA (ONNX) + KcELECTRA LoRA (ko) / DistilRoBERTa โ
+โ audio emotion (7-class) text emotion (7-class ko / 7-class en) โ
+โ โ
+โ fusion: per-language, per-class trained weights (v2) โ
+โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+ โ per-speaker emotion distribution
+ โผ
+โโโโโโโโโโโโโโโโโโโโโ Stage 3 (Youngkyun) โโโโโโโโโโโโโโโโโโโโ
+โ character_mapping + garden_logic + recap_generator โ
+โ 9 pair interactions 5 levels ยท 4 moods Claude LLM โ
+โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+ โ
+ โผ
+โโโโโโโโโโโโโโโโโโโโโ Stage 4 (Youngkyun) โโโโโโโโโโโโโโโโโโโโ
+โ FastAPI + SQLite โ React Native (Expo) โ
+โ 6 endpoints, async expo-router ยท SVG โ
+โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+```
+
+**End-to-end latency:** a 1-minute call finishes in about 2 minutes on the HF Spaces CPU deployment.
+**Live server:** `https://bbbakery-ustwo-api.hf.space`
+
+---
+
+## Demo โ the garden grows
+
+The emotion garden levels up as positive-ratio interactions accumulate. Level 1 shows just two seedlings. After many positive calls, Level 5 is a full bloom with flowers, trees, and creatures.
+
+
+ 
+ 
+
+
+| Level | Threshold (cumulative interactions) | Visual |
+|-------|-------------------------------------|--------|
+| 1 | 0โ2 | Two seedlings |
+| 2 | 3โ7 | Grass + small flowers |
+| 3 | 8โ14 | Trees + many flowers |
+| 4 | 15โ24 | Lush garden + creatures |
+| 5 | 25+ | Sunset sky + butterflies, rabbits + full bloom |
+
+A **mood overlay** (`happy` ยท `neglected` ยท `recovering` ยท `conflict`) is layered on top, decided by the recent positive/negative ratio and the days since the last interaction. A healthy garden that's been ignored for 5+ days turns `neglected`; a recent call with heavy negative emotion pushes it to `conflict`.
+
+---
+
+## Stage 1โ2 โ ML Pipeline (Seungjae)
+
+### Stage 1: Diarization + ASR
+
+| Component | Model | Config |
+|-----------|-------|--------|
+| VAD | Silero VAD | onset 0.5, min_speech 0.25s |
+| Diarization | pyannote-audio 4.x (3.1) | HF token required |
+| ASR | Faster-Whisper (large-v3-turbo) | INT8 quantized |
+| Forced alignment | WhisperX wav2vec2 | per-word timestamps |
+| Language ID | Whisper auto-detect + SenseVoice-Small | ko / en / unknown |
+
+**Output:** `Stage1Output` โ `segments: [{speaker, start, end, text, lang}]`, `processing_info`, `models`. Entry point: `src/stage1/process.py`.
+
+### Stage 2: Emotion Recognition (bilingual)
+
+| Channel | Model | Classes |
+|---------|-------|---------|
+| Audio emotion | `emotion2vec_plus_base` + **LoRA fine-tuned (ONNX)** | 7: neutral, joy, sadness, anger, surprise, fear, disgust |
+| Text emotion โ Korean | **KcELECTRA-base-v2022 + PEFT LoRA (ONNX)** | 7: neutral, joy, sadness, anger, surprise, fear, disgust |
+| Text emotion โ English | DistilRoBERTa (j-hartmann/emotion-english, zero-shot) | 7 Ekman + neutral |
+| Fusion | **Per-language, per-class weights trained via gradient descent** | EMOTION_FUSION_WEIGHTS_KO / _EN |
+
+### Evaluation โ English (RAVDESS, n=2,880)
+
+| Condition | Accuracy | Macro F1 | Latency |
+|-----------|----------|----------|---------|
+| Clean (studio) | **93.2%** | **0.932** | 250 ms |
+| Phone (PSTN simulation, 300โ3400 Hz band-limit) | **71.5%** | **0.710** | 242 ms |
+| Degradation | โ21.7 pp | โ0.222 | โ |
+
+- **Phone-robust emotions (F1 > 0.80):** anger (0.836), surprise (0.856) โ high-energy acoustic cues survive the band-limit.
+- **Phone-degraded (largest drop):** joy (0.946 โ 0.585, โ0.361), sadness (0.868 โ 0.559, โ0.309) โ subtler pitch/timbre cues degrade hardest.
+- **Text emotion sanity check (DistilRoBERTa on j-hartmann's held-out set):** 95.2% accuracy.
+
+Full report: [`docs/stage2/english-evaluation-report.md`](docs/stage2/english-evaluation-report.md)
+
+### Evaluation โ Korean (KcELECTRA fine-tuning)
+
+- Fine-tuned on the AI Hub *๊ฐ์ฑ ๋ํ ๋ง๋ญ์น* (emotion dialogue corpus) on Colab GPU.
+- Macro F1: **0.20 (base) โ 0.65 (fine-tuned)**, a 3.25ร improvement.
+- Discovered an unlabeled "joy" cluster in the dataset and applied class weighting to recover minority classes.
+
+### Fusion weight training (v2, 2026-04-19)
+
+Replaced greedy per-emotion grid search with **PyTorch gradient descent** training `w_a = sigmoid(ฮฑ)` per class, optimized jointly with cross-entropy + L2 regularization on a held-out val split.
+
+| Language | Training data | Samples | Val Macro F1 (v1 โ v2) |
+|----------|---------------|---------|------------------------|
+| English | JL-Corpus + SAVEE + MELD + RAVDESS phone | 2,821 | 0.6295 โ **0.7596 (+12.67%p)** |
+| Korean | AI Hub 263 val | 1,294 | 0.8736 โ **0.8748 (tie)** |
+
+Biggest single win: English `fear` F1 rose from **0.04 โ 0.67** (greedy's `audio_w=0.00` was a local trap). Per-language weight tables: `src/common/constants.py` (`EMOTION_FUSION_WEIGHTS_KO` / `_EN`). Details: [`docs/stage2/fusion-weights-english-grid-search.md`](docs/stage2/fusion-weights-english-grid-search.md).
+
+### End-to-end test โ MELD (Friends TV dialogue)
+
+We built 8 scenario WAVs from MELD (anger, joy, sadness, surprise, fear, bittersweet, annoyance, calm) and ran them through the deployed server.
+
+| Metric | Result |
+|--------|--------|
+| Pipeline success | **8 / 8** files completed Stage 1 โ 2 โ 3 |
+| Exact top-1 emotion label match | **7 / 7** (one tie) |
+| Average processing time | ~2 min / file (HF Spaces CPU) |
+
+### End-to-end test โ 20Hours Korean demo (2026-04-20)
+
+Companion Korean E2E set curated from the 20Hours Korean Conversational Speech dataset (M-F pairs only). 7 scenarios @ ~1 min each for live demo.
+
+| Metric | Result |
+|--------|--------|
+| Pipeline success | **7 / 7** files completed Stage 1 โ 2 โ 3 |
+| Intended-emotion match (one speaker) | **4 / 7** |
+| Source | `data/20hours_test/` + `scripts/test_20hours_e2e_server.py` |
+
+Note: 20Hours source is ASR-training data without emotion labels โ `ground_truth.json` lists *intended demo emotions* (for graph visualization), not ground-truth annotations.
+
+---
+
+## Stage 3โ4 โ App + Server (Youngkyun)
+
+### Stage 3: Reaction ยท Garden ยท Recap
+
+`src/stage3/process.py` takes the Stage 2 output and runs three independent modules:
+
+| Module | Role | Key logic |
+|--------|------|-----------|
+| `character_mapping.py` | Buckets each speaker's 7-class emotion into 4 moods (up / calm / down / tense) and looks up the pair cell in a 4ร4 matrix โ one of 9 pair states + a giver role | `joy โ up`, `anger โ tense`, `surprise/fear/disgust` resolved via residual distribution; `up ร down โ comforting (giver=A)`, `tense ร tense โ back_turned`, `calm ร calm โ idle`, ... |
+| `character_mapping.py` (intensity) | Emits a `CharacterReaction.intensity` in `{1, 2, 3}` used by the app for healing-cycle thresholds | Default 2 โ 3 cycles to heal; 1 โ 4 cycles; 3 โ 2 cycles |
+| `garden_logic.py` | Computes growth delta (0โ3) from call quality | `positive_ratio โฅ 0.5` โ +3, `pos โฅ 0.3 && neg < 0.3` โ +2, `neg โฅ 0.5` โ 0 |
+| `recap_generator.py` | Generates the narrative recap card via the Claude API | System prompt requires one concrete hook from the transcript (topic, decision, shared joke) + a light garden-voice framing โ titles are call-specific, not combo templates. Rule-based template fallback when no API key. |
+
+**Mood resolution:** `neg โฅ 0.5` โ `conflict`; `neg โฅ 0.3 && pos < 0.3` โ `recovering`; `level โค 1 && pos < 0.3` โ `neglected`; otherwise `happy`. **Level only goes up** โ tough calls shift the tint, never the count.
+
+**Confidence gate:** a speaker whose top emotion probability is below 0.5 is demoted to `calm` for pair-state lookup. The 7-class mood chip on the Results screen can therefore differ from the pair-state particle โ the chip shows the dominant label, the particle shows the matrix cell after gating.
+
+### Stage 4: FastAPI backend
+
+- **Framework:** FastAPI + SQLAlchemy + SQLite (local) ยท 4 tables: `calls`, `analysis_results`, `checkins`, `garden_state`.
+- **Async pipeline:** `POST /api/upload` โ `POST /api/analyze?call_id=X` returns 202 Accepted + a background thread, and the client polls `GET /api/calls/{id}`.
+- **Mock path:** drop `data/samples/{call_id}_stage2.json` and the API skips Stage 1โ2, running only Stage 3 โ useful for E2E testing without the heavy ML deps.
+- **Endpoints (6):** `/api/upload`, `/api/analyze`, `/api/calls`, `/api/calls/{id}`, `/api/checkins`, `/api/garden`. `/api/calls` extracts `recap_card.title` from the stored Stage 3 JSON so the Home and History feeds can headline each card with a distinct title.
+- **Tests:** 12 API tests on in-memory SQLite via pytest.
+
+### React Native (Expo) app
+
+**Router:** `expo-router` โ 3 tabs (`Us`, `History`, `Settings`) plus modal routes (`checkin`, `results/[callId]`).
+
+| Screen | Description |
+|--------|-------------|
+| 
| **Onboarding** โ paper-deck intro ("A garden for two") that sets the metaphor before the first call lands |
+| | **Home (`Us`)** โ live character scene + garden, recent-call feed, mailbox entry point |
+| 
| **Seed bloom alert** โ level-up moment; fires when the interaction count crosses a garden threshold, offering `View` / `Later` |
+| 
| **Check-in** โ 2-step prompt: my mood, then my guess for the partner's mood (empathic accuracy) |
+| 
| **Results โ emotion analysis** โ `My mood` / `Their mood` chips, suggestion line in the garden-voice, `Emotional Landscape` wave (uplifting โ heavy), and `Moments that mattered` per-speaker slices |
+| | **Results โ recap card** โ call-specific LLM title + narrative + highlights + per-call `Our Garden` delta, followed by the `Was this accurate?` thumbs-up/down feedback loop |
+| | **History** โ `Our Emotional Flow` chart (Me vs Partner), garden delta summary, recent moment cards headlined by the call-specific recap title |
+| 
| **Settings** โ language toggle (ko/en), developer mode, dev tools |
+
+**Character animation layers** (all built on `react-native-reanimated` + `react-native-svg`):
+
+1. **Idle breathing** โ per-mood rhythm profile (up / calm / down / tense), uniform `scaleAmp` 1.02, micro-bob, micro-sway (1.4s cycle).
+2. **Eye blink** โ 2.8โ5.5s interval, 15% chance of a double-blink.
+3. **Emotion transition** โ joy = bounce + arms **raised**, sadness = sink + drooping arms, anger / disgust = crossed arms (no tilt), surprise = pop, fear = shrink. Body rotation is globally forbidden โ emotion reads via face, pupil, body pose, and idle rhythm instead.
+4. **Pair-state body pose** โ `comforting` giver moves 60% closer with no tilt; receiver gets a gentle sag + `pupilOffset` (head-lowering effect on eyes).
+5. **Pair-state particle signature** โ one preset per 4ร4 matrix cell:
+ - `comforting` โ staged echo: 1 heart (giver โ receiver) + 1 bubble dot + 1 translucent heart (bezier engine, `useParticles`).
+ - `dancing` โ 1 music note rising above the heads.
+ - `cheering` / `listening` / `sitting_together` / `defusing` โ yellow sparkle cluster.
+ - `tension` / `back_turned` โ grey sigh cloud.
+ - `idle` (calm ร calm) โ silent; the couple just wanders.
+6. **Generalised healing** โ every giver-present cell (comforting, listening, defusing, cheering) counts interaction cycles; when the intensity-scaled threshold is hit the receiver transitions to neutral. If the giver's solo state was negative-caring (fear / anger / sadness / disgust), they heal too โ no one is left behind.
+7. **Tap interaction** โ spring bounce + floating hearts / sparkles (`FloatingHearts` component).
+8. **Wander** โ characters roam between pre-validated waypoints (3โ8s move, 2โ5s pause), with direction-aware facing (scaleX flip for left/right, back-view for upward movement).
+
+**Garden rendering (SVG):** Sky, Ground, Trees (level-specific types and counts), Flowers (distributed across four quadrants), Creatures (butterflies, rabbits, birds โ Level 4 and above).
+
+**i18n:** a custom `LocaleContext` drives Korean/English switching with `@ustwo/locale` persisted to AsyncStorage.
+
+### State persistence
+
+| Store | Data | Key |
+|-------|------|-----|
+| SQLite (server) | calls, analysis_results, checkins, garden_state | โ |
+| AsyncStorage (app) | locale, dev_mode, garden (interactionCount + lastPositiveRatio + lastInteractionDate), checkins (local cache) | `@ustwo/*` |
+
+---
+
+## Repository map
+
+```
+src/stage1/ Speaker diarization + ASR (Seungjae)
+src/stage2/ Audio + text emotion + fusion (Seungjae)
+src/stage3/ Character, garden, recap (Youngkyun)
+src/stage4/ FastAPI + SQLite + orchestration (Youngkyun)
+src/common/ Pydantic schemas shared across stages
+app/ React Native (Expo) app (Youngkyun)
+ src/routes/ expo-router (tabs, checkin, results, dev)
+ src/components/ characters, garden, scenes, layout
+ src/contexts/ Locale, Garden, DevMode
+ src/hooks/ Animation hooks (idle, blink, emotion, tap, wander)
+notebooks/ KcELECTRA fine-tuning (Colab)
+tests/ pytest (73 Python tests) + jest (71 JS tests)
+docs/stage{1,2,3,4}/ Per-stage technical docs
+docs/images/ README screenshots
+config.yaml Global config (model paths, thresholds)
+Dockerfile HF Spaces deployment (Python 3.12 + torch CPU + ffmpeg)
+```
+
+---
+
+## Getting started
+
+### 1. Clone
+
+```bash
+git clone https://github.com/boolooppang/UsTwo.git
+cd UsTwo
+```
+
+### 2. Backend (Python)
+
+```bash
+python -m venv venv && source venv/bin/activate
+pip install -r requirements.txt
+export HF_TOKEN=your_huggingface_token # required by pyannote
+
+uvicorn src.stage4.main:app --reload --port 8000
+# โ POST /api/upload, POST /api/analyze?call_id=X, GET /api/calls
+```
+
+### 3. App (React Native / Expo)
+
+```bash
+cd app && npm install
+npx expo start --dev-client
+# Press 'i' for iOS simulator, or scan the QR code with a real device
+```
+
+To point the app at a different server, edit `API_BASE` in `app/src/api/client.ts` (it currently defaults to the HF Spaces URL).
+
+### 4. Tests
+
+```bash
+python -m pytest tests/ -v # 73 Python tests
+cd app && npx jest # 71 JS tests (144 total)
+```
+
+---
+
+## Deployment โ HuggingFace Spaces (Docker)
+
+The API server is deployed to HuggingFace Spaces as a Docker image, running the full pipeline (Stage 1 โ 2 โ 3).
+
+- **Live URL:** `https://bbbakery-ustwo-api.hf.space`
+- **Environment variables:** `HF_TOKEN`, `ANTHROPIC_API_KEY` (set in Spaces Settings).
+- **Config files:** [`Dockerfile`](Dockerfile), [`railway.toml`](railway.toml), [`requirements-deploy.txt`](requirements-deploy.txt).
+
+```bash
+# Local Docker test
+docker build -t ustwo .
+docker run -p 7860:7860 -e HF_TOKEN=your_token ustwo
+```
+
+---
+
+## Team & ownership
+
+| Member | Role | Responsibility |
+|--------|------|----------------|
+| **Juhyun** | Product Owner / Researcher | Product design, UX research, empathic-accuracy literature review, emotion โ character mapping spec, evaluation plan, final paper |
+| **Seungjae** | ML Engineer | Stage 1โ2 end to end โ diarization, ASR, audio emotion, text emotion (ko/en), RAVDESS/MELD evaluation, KcELECTRA fine-tuning |
+| **Youngkyun** | App Engineer | Stage 3โ4 end to end โ character mapping implementation, garden logic, LLM recap, FastAPI + SQLite server, React Native app, HF Spaces Docker deployment |
+
+---
+
+## Success criteria
+
+**MVP โ achieved:**
+- โ
Upload a call recording โ character reaction + recap card within ~2 minutes.
+- โ
Happy vs tense calls produce visibly different character reactions (9 pair states).
+- โ
Cumulative positive calls grow the garden through 5 levels.
+- โ
Bilingual pipeline (Korean + English) with automatic language detection.
+- โ
MELD-based end-to-end test: 8/8 pipeline success, 7/7 exact emotion-label match.
+
+---
+
+## License
+
+MIT
diff --git a/VERSION b/VERSION
new file mode 100644
index 0000000000000000000000000000000000000000..227cea215648b1af34a87c9acf5b707fe02d2072
--- /dev/null
+++ b/VERSION
@@ -0,0 +1 @@
+2.0.0
diff --git a/config.yaml b/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..cdcf054a74dcb4391a7d7f4f78b0133894636039
--- /dev/null
+++ b/config.yaml
@@ -0,0 +1,87 @@
+# UsTwo Project Configuration
+project:
+ name: "UsTwo"
+ version: "0.1.0"
+
+paths:
+ data_dir: "data"
+ samples_dir: "data/samples"
+ models_dir: "data/models"
+ output_dir: "data"
+
+# Stage 1: Speaker Diarization + ASR
+stage1:
+ output_path: "data/stage1_output.json"
+ segments_dir: "data/segments"
+
+ preprocessing:
+ target_sample_rate: 16000
+ max_duration_sec: 300
+ min_duration_sec: 3
+ target_peak: 0.95 # peak normalization โ handles volume differences across devices
+
+ diarization:
+ model: "pyannote/speaker-diarization-3.1"
+ num_speakers: 2
+ merge_gap_sec: 0.15
+
+ asr:
+ model: "large-v3-turbo"
+ compute_type: "int8"
+ language: null # null = auto-detect
+ batch_size: 16
+
+ alignment:
+ enabled: true
+
+ language_id:
+ enabled: true
+ # SenseVoice disabled โ Whisper language + text heuristic ์ฌ์ฉ
+ # emotion2vec finetuning ์คํจ ์ SenseVoice ๊ฐ์ ํํธ ํ์ฉ ์์
+ # model: "FunAudioLLM/SenseVoiceSmall"
+
+# Stage 2: Audio Emotion + Text Emotion
+stage2:
+ input_path: "data/stage1_output.json"
+ output_path: "data/stage2_output.json"
+ audio_emotion:
+ model: "iic/emotion2vec_plus_base"
+ lora_onnx_path: "data/models/lora_emotion2vec_7class/model.onnx"
+ finetuned_checkpoint: null # legacy, use lora_onnx_path instead
+ text_emotion:
+ korean_model: "searle-j/kote_for_easygoing_people"
+ korean_lora_onnx_path: "data/models/lora_kcelectra_7class/model.onnx"
+ korean_lora_tokenizer: "data/models/lora_kcelectra_7class/best_model"
+ english_model: "j-hartmann/emotion-english-distilroberta-base"
+ fusion:
+ mode: "emotion_specific" # "emotion_specific" (per-class grid-search optimized) or "fixed" (60/40)
+ audio_weight: 0.6 # fallback for mode="fixed"
+ text_weight: 0.4
+
+# Stage 3: Character Reaction + Garden + Recap
+stage3:
+ input_path: "data/stage2_output.json"
+ output_path: "data/stage3_output.json"
+ recap:
+ llm_provider: "anthropic" # anthropic | openai
+ model: "claude-sonnet-4-20250514"
+ max_tokens: 500
+ garden:
+ growth_per_call: 5
+ max_level: 5
+
+# Stage 4: FastAPI Server
+stage4:
+ input_path: "data/stage3_output.json"
+ host: "0.0.0.0"
+ port: 8000
+ max_upload_size_mb: 50
+ allowed_extensions: [".wav", ".mp3", ".m4a", ".ogg"]
+
+# API Keys (ํ๊ฒฝ ๋ณ์์์ ๋ก๋ โ .env ํ์ผ ์ฌ์ฉ)
+api: {}
+
+# Logging
+logging:
+ level: "INFO"
+ format: "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
diff --git a/data/meld_test/01_angry_fight.wav b/data/meld_test/01_angry_fight.wav
new file mode 100644
index 0000000000000000000000000000000000000000..5467df50b9bbe2c906c30bedd34b88db83d12fb7
--- /dev/null
+++ b/data/meld_test/01_angry_fight.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a9e440869de56c29b94e85b12b6145181439011f1c2a24831ee58b17d268bda7
+size 1098486
diff --git a/data/meld_test/02_happy_loving.wav b/data/meld_test/02_happy_loving.wav
new file mode 100644
index 0000000000000000000000000000000000000000..ec732cdb912432220f6a4d2e9dd5bbfe88b3a1fa
--- /dev/null
+++ b/data/meld_test/02_happy_loving.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:194ce7561f2bb9efa54386938fb2fc0d87713590cbb7e728d7d3185c19f101cd
+size 1400224
diff --git a/data/meld_test/03_sad_emotional.wav b/data/meld_test/03_sad_emotional.wav
new file mode 100644
index 0000000000000000000000000000000000000000..83d3a96dd9763d400e9283be91c501fa30afa0c0
--- /dev/null
+++ b/data/meld_test/03_sad_emotional.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:23c3e6458334dc0b24940144b9e202c82b7b86be6c0cdea982a76e7934206848
+size 1584548
diff --git a/data/meld_test/04_surprise_shock.wav b/data/meld_test/04_surprise_shock.wav
new file mode 100644
index 0000000000000000000000000000000000000000..4cd9baa13e085de70871b8a2869d70117f01d8e4
--- /dev/null
+++ b/data/meld_test/04_surprise_shock.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2958c4ad34b23c793995b9357f80e5a3a89cf26d8675eff581402b17a2c01e9f
+size 966734
diff --git a/data/meld_test/05_fear_anxiety.wav b/data/meld_test/05_fear_anxiety.wav
new file mode 100644
index 0000000000000000000000000000000000000000..6d993cd2cbc58b7b5a49c239968c5b8dc86373a1
--- /dev/null
+++ b/data/meld_test/05_fear_anxiety.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3fdad57492923cc788a550f75a486950e61d9b851cf69a2e9a694da5c14bca8a
+size 637008
diff --git a/data/meld_test/06_disgust_annoyance.wav b/data/meld_test/06_disgust_annoyance.wav
new file mode 100644
index 0000000000000000000000000000000000000000..73822d233bd49297045c153dceb44e6056b345dd
--- /dev/null
+++ b/data/meld_test/06_disgust_annoyance.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5dd1600ef131cfe92ba93f79a5d82396164aade9902d8446ebef475480a2cb6
+size 933964
diff --git a/data/meld_test/07_bittersweet.wav b/data/meld_test/07_bittersweet.wav
new file mode 100644
index 0000000000000000000000000000000000000000..842b4024bbd4730650a802f8ebd0ba1f11b96ff5
--- /dev/null
+++ b/data/meld_test/07_bittersweet.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:11c445fc5445a7e9c618b09deb351852ac616580ff7d59ac579a0be980a059f3
+size 1384526
diff --git a/data/meld_test/08_calm_daily.wav b/data/meld_test/08_calm_daily.wav
new file mode 100644
index 0000000000000000000000000000000000000000..f504bc2851fe400d2dc79753bbe67c58df060a64
--- /dev/null
+++ b/data/meld_test/08_calm_daily.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:05b2cc7fdac461bf9e8d6610388891552806a6a1350c55b4624f8209dcc7ca07
+size 1294414
diff --git a/data/meld_test/09_opposite_emotions.wav b/data/meld_test/09_opposite_emotions.wav
new file mode 100644
index 0000000000000000000000000000000000000000..d93a252b08b3ff85d37580c99b46b981e90640c1
--- /dev/null
+++ b/data/meld_test/09_opposite_emotions.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ce64da9783391134b20071c315ff794d9508b14c8c8d639a80e6d09267f1bb4f
+size 796766
diff --git a/data/meld_test/README.md b/data/meld_test/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..5693e22282fae66b9e660dc7a6b26f35f919e5e9
--- /dev/null
+++ b/data/meld_test/README.md
@@ -0,0 +1,53 @@
+# MELD English Test Sets
+
+## Emotion Label Alignment
+
+| UsTwo Pipeline (EN) | MELD Label | Match |
+|---|---|---|
+| neutral | neutral | โ
Exact |
+| joy | joy | โ
Exact |
+| sadness | sadness | โ
Exact |
+| anger | anger | โ
Exact |
+| surprise | surprise | โ
Exact |
+| fear | fear | โ
Exact |
+| disgust | disgust | โ
Exact |
+
+**7/7 labels match exactly.** No mapping needed.
+
+## Test Sets
+
+| File | Scenario | Speakers | Primary Emotion | Duration | Utterances | Emotion Distribution |
+|---|---|---|---|---|---|---|
+| 01_angry_fight | Couple in a heated argument | Ross, Rachel | anger | 36.9s | 11 utts | anger:7 neutral:2 sadness:1 disgust:1 |
+| 02_happy_loving | Couple being affectionate and playful | Chandler, Monica | joy | 43.8s | 16 utts | joy:6 surprise:5 anger:3 neutral:1 sadness:1 |
+| 03_sad_emotional | Emotional confession โ "you still love me?" | Ross, Rachel | sadness | 55.4s | 17 utts | neutral:7 sadness:4 anger:3 surprise:3 |
+| 04_surprise_shock | Drunk voicemail surprise scene | Ross, Rachel | surprise | 30.2s | 8 utts | surprise:5 neutral:2 sadness:1 |
+| 05_fear_anxiety | Anxious and worried conversation | Chandler, Rachel | fear | 19.9s | 12 utts | fear:5 neutral:4 surprise:2 sadness:1 |
+| 06_disgust_annoyance | Annoyed and frustrated bickering | Joey, Rachel | anger | 29.2s | 11 utts | anger:5 neutral:2 sadness:1 surprise:1 fear:1 joy:1 |
+| 07_bittersweet | Saying goodbye with conflicting feelings | Ross, Rachel | sadness | 43.3s | 14 utts | sadness:6 surprise:3 anger:3 fear:1 neutral:1 |
+| 08_calm_daily | Casual everyday chitchat (baseline) | Joey, Monica | neutral | 40.4s | 15 utts | neutral:13 joy:2 |
+| 09_opposite_emotions | Tense speaker + calm listener โ triggers `listening` pair animation | 2 spk | surprise | 24.9s | 11 segs | surprise:5 neutral:3 anger:1 fear:1 joy:1 (pipeline fused) |
+
+## Notes
+- All dialogues are 2-speaker (male + female) conversations
+- 03: Ross+Rachel only, utt15 removed (timestamp overlap with utt14)
+- 04: starts from "Rach, I got a message from you", utt3/utt8 removed (timestamp overlaps)
+- 06: Joey+Rachel only, utt16 removed (addresses Ross)
+- 07: utt11 removed (timestamp overlap with utt10)
+- 09: **opposite-emotion demo scene (updated 2026-04-22)** โ WAV replaced by user. E2E pipeline output: speaker_0 dominant=surprise (max 0.68, tense via anger residual), speaker_1 dominant=neutral (max 0.54, calm). `(tense, calm)` โ **pair_state=`listening`** triggers (sparkles effect, speaker_1 = giver/listener). Recap captures "one expressing worry through sharp words, the other sharing tender vision." Per-utterance ground-truth labels not available post-update; distribution numbers derived from pipeline fused output.
+
+## Source
+- Dataset: MELD (Multimodal EmotionLines Dataset)
+- Source: Friends TV series
+- Paper: Poria et al., ACL 2019
+- Each WAV is a full dialogue concatenated from per-utterance MP4 clips
+- Audio: 16kHz mono PCM (matches pipeline input format)
+
+## Usage
+```bash
+# Run pipeline on a single test set
+python scripts/run_pipeline.py data/meld_test/01_angry_fight.wav
+
+# Evaluate all test sets
+python scripts/evaluate_meld_test.py
+```
diff --git a/data/meld_test/ground_truth.json b/data/meld_test/ground_truth.json
new file mode 100644
index 0000000000000000000000000000000000000000..4a266b1a9272e6da88976d9f1bd21e1b504fcd16
--- /dev/null
+++ b/data/meld_test/ground_truth.json
@@ -0,0 +1,772 @@
+{
+ "01_angry_fight": {
+ "description": "Ross-Rachel breakup fight โ anger dominant (S3E15)",
+ "scenario": "Couple in a heated argument",
+ "primary_emotion": "anger",
+ "source": "MELD Friends S3E15 Dialogue 51 (utt0 removed, overlap fix)",
+ "duration_sec": 34.3,
+ "emotion_distribution": {
+ "sadness": 1,
+ "neutral": 1,
+ "disgust": 1,
+ "anger": 7
+ },
+ "total_utterances": 10,
+ "utterances": [
+ {
+ "speaker": "Rachel",
+ "emotion": "sadness",
+ "sentiment": "negative",
+ "text": "Hi. Look um, about what happened earlier..."
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "No, hey, well, I-I completely understand. You were, you were stressed."
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "disgust",
+ "sentiment": "negative",
+ "text": "I was gonna give you a chance to apologise to me."
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "For what? For letting you throw me out of your office?"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "You had no"
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "Yeah, well excuse me for wanting to be with my girlfriend on our anniversary, boy what an ass am I."
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "But I told you, I didnโt have the time!"
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "Yeah, well you never have the time. I mean, I donโt feel like I even have a girlfriend anymore, Rachel."
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "Wh, Ross what do you want from me?"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "You want me, you want me to quit my job so you can feel like you have a girlfriend?"
+ }
+ ]
+ },
+ "02_happy_loving": {
+ "description": "Monica-Chandler sweet moment โ joy dominant (S5E14)",
+ "scenario": "Couple being affectionate and playful",
+ "primary_emotion": "joy",
+ "source": "MELD Friends S5E14 Dialogue 1026",
+ "duration_sec": 43.8,
+ "emotion_distribution": {
+ "joy": 6,
+ "neutral": 1,
+ "surprise": 5,
+ "anger": 3,
+ "sadness": 1
+ },
+ "total_utterances": 16,
+ "utterances": [
+ {
+ "speaker": "Monica",
+ "emotion": "joy",
+ "sentiment": "positive",
+ "text": "You are so cute! How did you get to be so cute?"
+ },
+ {
+ "speaker": "Chandler",
+ "emotion": "joy",
+ "sentiment": "positive",
+ "text": "Well, my Grandfather was Swedish and my Grandmother was actually a tiny little bunny."
+ },
+ {
+ "speaker": "Monica",
+ "emotion": "joy",
+ "sentiment": "positive",
+ "text": "Okay, now you're even cuter!!"
+ },
+ {
+ "speaker": "Chandler",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "Y'know that is a popular opinion today I must say."
+ },
+ {
+ "speaker": "Monica",
+ "emotion": "surprise",
+ "sentiment": "negative",
+ "text": "What?"
+ },
+ {
+ "speaker": "Chandler",
+ "emotion": "surprise",
+ "sentiment": "negative",
+ "text": "The weirdest thing happened at the coffee house, I think, I think Phoebe was hitting on me."
+ },
+ {
+ "speaker": "Monica",
+ "emotion": "surprise",
+ "sentiment": "negative",
+ "text": "What are you talking about?"
+ },
+ {
+ "speaker": "Chandler",
+ "emotion": "joy",
+ "sentiment": "positive",
+ "text": "I'm telling you I think Phoebe thinks I'm foxy."
+ },
+ {
+ "speaker": "Monica",
+ "emotion": "joy",
+ "sentiment": "positive",
+ "text": "That's not possible!"
+ },
+ {
+ "speaker": "Chandler",
+ "emotion": "surprise",
+ "sentiment": "positive",
+ "text": "Ow!"
+ },
+ {
+ "speaker": "Monica",
+ "emotion": "joy",
+ "sentiment": "positive",
+ "text": "I'm sorry it's just, Phoebe just always thought you were, you were charming in a, in a sexless kind of way."
+ },
+ {
+ "speaker": "Chandler",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "Oh, y'know I-I can't hear that enough."
+ },
+ {
+ "speaker": "Monica",
+ "emotion": "sadness",
+ "sentiment": "negative",
+ "text": "I'm sorry, I think that you just misunderstood her."
+ },
+ {
+ "speaker": "Chandler",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "No, I didn't misunderstand, okay? She was all over me! She touched my bicep for crying out loud!"
+ },
+ {
+ "speaker": "Monica",
+ "emotion": "surprise",
+ "sentiment": "negative",
+ "text": "This bicep?"
+ },
+ {
+ "speaker": "Chandler",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "Well it's not flexed right now!"
+ }
+ ]
+ },
+ "03_sad_emotional": {
+ "description": "Ross-Rachel emotional confession โ sadness dominant (S3E25)",
+ "scenario": "Emotional conversation with sadness and regret",
+ "primary_emotion": "sadness",
+ "source": "MELD Friends S3E25 Dialogue 312 (Ross+Rachel only, utt15/19 removed)",
+ "duration_sec": 49.5,
+ "emotion_distribution": {
+ "neutral": 6,
+ "anger": 3,
+ "sadness": 4,
+ "surprise": 3
+ },
+ "total_utterances": 16,
+ "utterances": [
+ {
+ "speaker": "Ross",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "You donโt know?! Rach, you balded my girlfriend!"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "All right! Ross, do you think itโs easy for me to see you with somebody else?"
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "Y'know, hey! Youโre the one who ended it, remember?"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "Yeah, because I was"
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "sadness",
+ "sentiment": "negative",
+ "text": "You still love me?"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "sadness",
+ "sentiment": "negative",
+ "text": "Noo."
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "You still love me."
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "surprise",
+ "sentiment": "positive",
+ "text": "Oh, y-yeah, so, you-you love me!"
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "surprise",
+ "sentiment": "positive",
+ "text": "Noo, nnnnn. What does this mean? What do you, I mean do you wanna, get back together?"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "Noo!"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "surprise",
+ "sentiment": "positive",
+ "text": "Maybe!"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "I, I donโt know."
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "sadness",
+ "sentiment": "negative",
+ "text": "Ross, I still canโt forgive you for what you did, I canโt, I just, but sometimes when Iโm with you I just, I feel so..."
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "sadness",
+ "sentiment": "negative",
+ "text": "What?!"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "I just, I feel, I-I just..."
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "I feel..."
+ }
+ ]
+ },
+ "04_surprise_shock": {
+ "description": "Ross-Rachel drunk voicemail surprise (S4E21)",
+ "scenario": "Unexpected news and surprising revelations between couple",
+ "primary_emotion": "surprise",
+ "source": "MELD Friends S4E21 Dialogue 848 (from utt2, overlaps removed)",
+ "duration_sec": 30.2,
+ "emotion_distribution": {
+ "neutral": 2,
+ "sadness": 1,
+ "surprise": 5
+ },
+ "total_utterances": 8,
+ "utterances": [
+ {
+ "speaker": "Ross",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "Rach, I got a message from you."
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "sadness",
+ "sentiment": "negative",
+ "text": "Oh my God Ross, no, hang up the phone, give me the phone Ross, give me the phone, give me the phone, give me the."
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "surprise",
+ "sentiment": "negative",
+ "text": "You're over me?"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "surprise",
+ "sentiment": "negative",
+ "text": "Ohhhhhhhh God."
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "surprise",
+ "sentiment": "negative",
+ "text": "Wha... you're uh, you're, you're over me?"
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "surprise",
+ "sentiment": "negative",
+ "text": "When, when were you... under me?"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "Well, basically, lately, I've uh, I've uh, sort of had feelings for you."
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "surprise",
+ "sentiment": "negative",
+ "text": "OK, I need to lie down."
+ }
+ ]
+ },
+ "05_fear_anxiety": {
+ "description": "Chandler-Rachel anxious situation โ fear dominant (S7E11)",
+ "scenario": "Anxious and worried conversation between two people",
+ "primary_emotion": "fear",
+ "source": "MELD Friends S7E11 Dialogue 989",
+ "duration_sec": 19.9,
+ "emotion_distribution": {
+ "surprise": 2,
+ "neutral": 4,
+ "sadness": 1,
+ "fear": 5
+ },
+ "total_utterances": 12,
+ "utterances": [
+ {
+ "speaker": "Rachel",
+ "emotion": "surprise",
+ "sentiment": "positive",
+ "text": "Its still there!"
+ },
+ {
+ "speaker": "Chandler",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "Mrs. Braverman must be out."
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "sadness",
+ "sentiment": "negative",
+ "text": "She could be out of town. Maybe sheโll be gone for months."
+ },
+ {
+ "speaker": "Chandler",
+ "emotion": "fear",
+ "sentiment": "negative",
+ "text": "By then, the cheesecake may have gone bad. We donโt want her to come back to bad cheesecake."
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "fear",
+ "sentiment": "negative",
+ "text": "No that could kill her."
+ },
+ {
+ "speaker": "Chandler",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "Well, we donโt want that."
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "No, so weโre protecting her."
+ },
+ {
+ "speaker": "Chandler",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "But we should take it."
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "fear",
+ "sentiment": "negative",
+ "text": "But we should move quick."
+ },
+ {
+ "speaker": "Chandler",
+ "emotion": "surprise",
+ "sentiment": "negative",
+ "text": "Why?"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "fear",
+ "sentiment": "negative",
+ "text": "Because I think I just heard her moving around in there."
+ },
+ {
+ "speaker": "Chandler",
+ "emotion": "fear",
+ "sentiment": "negative",
+ "text": "Go! Go! Go! Go! Go! Go! Go! Go! Go! Go!"
+ }
+ ]
+ },
+ "06_disgust_annoyance": {
+ "description": "Joey-Rachel annoyance/bickering scene (S6E9)",
+ "scenario": "Annoyed and frustrated reactions between couple",
+ "primary_emotion": "anger",
+ "source": "MELD Friends S6E9 Dialogue 1025 (Joey+Rachel only, utt16 removed)",
+ "duration_sec": 29.2,
+ "emotion_distribution": {
+ "sadness": 1,
+ "surprise": 1,
+ "anger": 5,
+ "neutral": 2,
+ "fear": 1,
+ "joy": 1
+ },
+ "total_utterances": 11,
+ "utterances": [
+ {
+ "speaker": "Joey",
+ "emotion": "sadness",
+ "sentiment": "negative",
+ "text": "Will you hurry up?"
+ },
+ {
+ "speaker": "Joey",
+ "emotion": "surprise",
+ "sentiment": "negative",
+ "text": "Did you not hear me before when I told you that all of Janineโs friends are dancers?!"
+ },
+ {
+ "speaker": "Joey",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "And that theyโre going to be drinking alot!"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "No, I did, but tell me again, because itโs so romantic."
+ },
+ {
+ "speaker": "Joey",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "Well youโre whippinโ so slow! Canโt you do it any faster?"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "Joey!"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "Come on!"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "I donโt wanna make any mistakes, alright?"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "fear",
+ "sentiment": "negative",
+ "text": "This is the only dessert and if I screw it up everybody's gonna be like โOh, remember that Thanksgiving when Rachel screwed up the trifle?โ"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "So why donโt you just let me worry about making the trifle and you just worry about eating it, alright?"
+ },
+ {
+ "speaker": "Joey",
+ "emotion": "joy",
+ "sentiment": "positive",
+ "text": "Oh I am!"
+ }
+ ]
+ },
+ "07_bittersweet": {
+ "description": "Ross-Rachel bittersweet farewell โ sadness+surprise (S5E5), overlap fixed",
+ "scenario": "Mixed emotions: saying goodbye with conflicting feelings",
+ "primary_emotion": "sadness",
+ "source": "MELD Friends S5E5 Dialogue 676 (utt11 overlap removed)",
+ "duration_sec": 43.3,
+ "emotion_distribution": {
+ "sadness": 6,
+ "surprise": 3,
+ "fear": 1,
+ "anger": 3,
+ "neutral": 1
+ },
+ "total_utterances": 14,
+ "utterances": [
+ {
+ "speaker": "Ross",
+ "emotion": "sadness",
+ "sentiment": "negative",
+ "text": "is for me not to see you anymore."
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "surprise",
+ "sentiment": "positive",
+ "text": "That's crazy!"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "surprise",
+ "sentiment": "positive",
+ "text": "You can't do that!"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "sadness",
+ "sentiment": "negative",
+ "text": "What are you going to tell her?"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "fear",
+ "sentiment": "negative",
+ "text": "Oh God."
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "sadness",
+ "sentiment": "negative",
+ "text": "Ohh, you already agreed to this, haven't you?"
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "sadness",
+ "sentiment": "negative",
+ "text": "It's awful I know, I mean, I feel terrible but I have to do this if I want my marriage to work."
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "sadness",
+ "sentiment": "negative",
+ "text": "And I do, I have to make"
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "surprise",
+ "sentiment": "positive",
+ "text": "Ohh! Lucky me! Oh my God! That"
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "sadness",
+ "sentiment": "negative",
+ "text": "You have no idea what a nightmare this has been. This is so hard."
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "Oh yeah, really? Is it Ross? Yeah? Okay, well let me make this a just a little bit easier for you."
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "Storming out!"
+ },
+ {
+ "speaker": "Ross",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "Rachel, this is your apartment."
+ },
+ {
+ "speaker": "Rachel",
+ "emotion": "anger",
+ "sentiment": "negative",
+ "text": "Yeah, well that's how mad I am!!"
+ }
+ ]
+ },
+ "08_calm_daily": {
+ "description": "Joey-Monica casual conversation โ neutral dominant (S7E19)",
+ "scenario": "Normal everyday chitchat between friends (baseline)",
+ "primary_emotion": "neutral",
+ "source": "MELD Friends S7E19 Dialogue 8 (dev)",
+ "duration_sec": 40.4,
+ "emotion_distribution": {
+ "neutral": 13,
+ "joy": 2
+ },
+ "total_utterances": 15,
+ "utterances": [
+ {
+ "speaker": "Monica",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "Hey! What did you decide to do about the movie?"
+ },
+ {
+ "speaker": "Joey",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "I donโt know!"
+ },
+ {
+ "speaker": "Joey",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "Itโs not like itโs porn!"
+ },
+ {
+ "speaker": "Joey",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "This is a serious, legitimate movie."
+ },
+ {
+ "speaker": "Joey",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "And the nudity is really important to the story."
+ },
+ {
+ "speaker": "Monica",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "Thatโs what you say about porn."
+ },
+ {
+ "speaker": "Joey",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "Youโre right. Maybe I shouldnโt even go on the call back."
+ },
+ {
+ "speaker": "Monica",
+ "emotion": "joy",
+ "sentiment": "positive",
+ "text": "No! No you should! A lot of major actors do nude scenes! I mean, the chance to star in a movie? Come on!"
+ },
+ {
+ "speaker": "Joey",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "Well thatโs true."
+ },
+ {
+ "speaker": "Joey",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "And I am only naked in one scene."
+ },
+ {
+ "speaker": "Joey",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "Plus it sounds really great."
+ },
+ {
+ "speaker": "Joey",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "My characterโs catholic and he falls in love with this Jewish girl."
+ },
+ {
+ "speaker": "Joey",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "Who run away together and they get caught in this big rainstorm."
+ },
+ {
+ "speaker": "Joey",
+ "emotion": "neutral",
+ "sentiment": "neutral",
+ "text": "So we go into this barn and undress each other and hold each other."
+ },
+ {
+ "speaker": "Joey",
+ "emotion": "joy",
+ "sentiment": "positive",
+ "text": "Itโs really sweet and-and tender."
+ }
+ ]
+ },
+ "09_opposite_emotions": {
+ "description": "Updated clip โ pair_state 'listening' triggers (tense speaker + calm listener)",
+ "scenario": "One speaker worked-up/tense + one calm listener โ triggers listening pair animation (sparkles effect)",
+ "primary_emotion": "surprise",
+ "source": "MELD Friends clip (user-updated 2026-04-22, duration 24.9s)",
+ "duration_sec": 24.9,
+ "pair_state_observed": "listening",
+ "emotion_distribution": {
+ "surprise": 5,
+ "anger": 1,
+ "neutral": 3,
+ "fear": 1,
+ "joy": 1
+ },
+ "per_speaker_fused": {
+ "speaker_0": {
+ "surprise": 4,
+ "anger": 1
+ },
+ "speaker_1": {
+ "surprise": 1,
+ "neutral": 3,
+ "fear": 1,
+ "joy": 1
+ }
+ },
+ "total_segments": 11,
+ "notes": "Ground-truth utterance labels not available after WAV update. Distribution derived from pipeline fused output."
+ }
+}
\ No newline at end of file
diff --git a/railway.toml b/railway.toml
new file mode 100644
index 0000000000000000000000000000000000000000..8b107a9b7c789e6b1c302f5222dee7c647a50343
--- /dev/null
+++ b/railway.toml
@@ -0,0 +1,9 @@
+[build]
+builder = "dockerfile"
+
+[deploy]
+startCommand = "uvicorn src.stage4.main:app --host 0.0.0.0 --port $PORT"
+healthcheckPath = "/api/health"
+healthcheckTimeout = 300
+restartPolicyType = "on_failure"
+restartPolicyMaxRetries = 3
diff --git a/requirements-deploy.txt b/requirements-deploy.txt
new file mode 100644
index 0000000000000000000000000000000000000000..076b592a0a79c2b7e504edd4683075e5a9ac755d
--- /dev/null
+++ b/requirements-deploy.txt
@@ -0,0 +1,31 @@
+# UsTwo Railway Deployment Dependencies
+# Stage 1 (pyannote, whisperx) ์ ์ธ โ ์์ด Stage 2 + Stage 3 + Stage 4 only
+# torch๋ Dockerfile์์ CPU-only๋ก ๋ณ๋ ์ค์น
+
+# Stage 4: FastAPI server
+fastapi>=0.100.0
+uvicorn>=0.20.0
+pydantic>=2.0.0
+sqlalchemy>=2.0.0
+python-multipart>=0.0.6
+pyyaml>=6.0
+
+# Stage 3: Recap generation
+anthropic>=0.20.0
+
+# Quick Recap: Whisper API transcription
+openai>=1.12.0
+
+# Stage 1: Speaker diarization + ASR
+pyannote.audio>=3.1
+faster-whisper>=1.0.0
+whisperx>=3.1.0
+torchaudio>=2.0.0
+
+# Stage 2: Emotion analysis (English)
+transformers>=4.38.0
+funasr>=1.0.0
+onnxruntime>=1.17.0
+librosa>=0.10.0
+soundfile>=0.12.0
+scipy>=1.10.0
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..802819f4908eb6fb77b16ece68ffbb6222cbc53c
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,17 @@
+fastapi>=0.100.0
+uvicorn>=0.20.0
+pydantic>=2.0.0
+sqlalchemy>=2.0.0
+python-multipart>=0.0.6
+pyyaml>=6.0
+anthropic>=0.20.0
+
+# ML โ Stage 2 emotion analysis
+torch>=2.0.0
+transformers>=4.38.0
+funasr>=1.0.0
+onnxruntime>=1.17.0
+librosa>=0.10.0
+soundfile>=0.12.0
+scikit-learn>=1.3.0
+scipy>=1.10.0
diff --git a/scripts/add_ravdess_to_english_manifest.py b/scripts/add_ravdess_to_english_manifest.py
new file mode 100644
index 0000000000000000000000000000000000000000..c3baa5a428c59dcdf447f65e37e25b066c5c5ebf
--- /dev/null
+++ b/scripts/add_ravdess_to_english_manifest.py
@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+"""Append RAVDESS fear/disgust/sadness (phone) to the English fusion manifest.
+
+Input:
+ - data/english_fusion/manifest.json (1,669 samples)
+ - data/ravdess/manifest.csv (RAVDESS phone/clean paths)
+
+Output:
+ - data/english_fusion/manifest_v2.json (2,821 samples)
+
+RAVDESS statement โ text:
+ 1 โ "Kids are talking by the door."
+ 2 โ "Dogs are sitting by the door."
+"""
+from __future__ import annotations
+
+import csv
+import json
+import logging
+from collections import Counter
+from pathlib import Path
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+RAVDESS_TARGETS = {"fear", "disgust", "sadness"}
+STATEMENT_TEXT = {
+ "1": "Kids are talking by the door.",
+ "2": "Dogs are sitting by the door.",
+}
+
+BASE_MANIFEST = Path("data/english_fusion/manifest.json")
+RAVDESS_CSV = Path("data/ravdess/manifest.csv")
+OUT = Path("data/english_fusion/manifest_v2.json")
+
+
+def load_ravdess_rows() -> list[dict]:
+ rows: list[dict] = []
+ with open(RAVDESS_CSV, newline="") as f:
+ reader = csv.DictReader(f)
+ for r in reader:
+ if r["emotion"] not in RAVDESS_TARGETS:
+ continue
+ text = STATEMENT_TEXT.get(r["statement"])
+ if not text:
+ continue
+ rows.append({
+ "path": r["phone_path"],
+ "text": text,
+ "label": r["emotion"],
+ "source": "ravdess_phone",
+ "speaker": f"ravdess_actor_{int(r['actor_id']):02d}",
+ })
+ return rows
+
+
+def main() -> None:
+ base = json.loads(BASE_MANIFEST.read_text())
+ logger.info("Base manifest: %d samples", len(base))
+
+ rav = load_ravdess_rows()
+ logger.info("RAVDESS additions: %d samples (fear/disgust/sadness)", len(rav))
+
+ combined = base + rav
+ counts = Counter(r["label"] for r in combined)
+ sources = Counter(r["source"] for r in combined)
+ logger.info("Total v2: %d", len(combined))
+ logger.info("By label: %s", dict(counts))
+ logger.info("By source: %s", dict(sources))
+
+ OUT.write_text(json.dumps(combined, indent=2, ensure_ascii=False))
+ logger.info("Saved to %s", OUT)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/asr_savee_disgust_surprise.py b/scripts/asr_savee_disgust_surprise.py
new file mode 100644
index 0000000000000000000000000000000000000000..3c5918e370096b1efbd506c594f42daf095b6fcb
--- /dev/null
+++ b/scripts/asr_savee_disgust_surprise.py
@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+"""Run faster-whisper ASR on SAVEE disgust + surprise wavs (120 files).
+
+Output: data/savee/savee_asr.json โ {wav_name: transcript}
+"""
+from __future__ import annotations
+
+import json
+import logging
+import re
+from pathlib import Path
+
+from faster_whisper import WhisperModel
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+SAVEE_DIR = Path("data/savee/ALL")
+OUT = Path("data/savee/savee_asr.json")
+PATTERN = re.compile(r"^[A-Z]{2}_(d|su)\d+\.wav$")
+
+
+def main() -> None:
+ wavs = sorted([p for p in SAVEE_DIR.iterdir() if PATTERN.match(p.name)])
+ logger.info("Found %d SAVEE disgust/surprise wavs", len(wavs))
+
+ logger.info("Loading faster-whisper large-v3-turbo (int8, CPU)...")
+ model = WhisperModel("large-v3-turbo", device="cpu", compute_type="int8")
+
+ results: dict[str, str] = {}
+ if OUT.exists():
+ results = json.loads(OUT.read_text())
+ logger.info("Loaded %d cached transcripts", len(results))
+
+ for i, wav in enumerate(wavs):
+ if wav.name in results:
+ continue
+ segments, _ = model.transcribe(str(wav), language="en", beam_size=5, vad_filter=False)
+ text = " ".join(s.text.strip() for s in segments).strip()
+ results[wav.name] = text
+ if i % 10 == 0:
+ OUT.write_text(json.dumps(results, indent=2, ensure_ascii=False))
+ logger.info("[%d/%d] %s -> %s", i + 1, len(wavs), wav.name, text[:60])
+
+ OUT.write_text(json.dumps(results, indent=2, ensure_ascii=False))
+ logger.info("Saved %d transcripts to %s", len(results), OUT)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/benchmark_emotion2vec.py b/scripts/benchmark_emotion2vec.py
new file mode 100644
index 0000000000000000000000000000000000000000..1a91569ebb6ceb17860901cd68696b8cd745a12e
--- /dev/null
+++ b/scripts/benchmark_emotion2vec.py
@@ -0,0 +1,513 @@
+#!/usr/bin/env python3
+"""emotion2vec Variant Smoke Test โ base vs plus_base vs plus_large ์ค์ธก ๋น๊ต.
+
+๊ธฐ์กด Stage 1 ์ถ๋ ฅ ์ธ๊ทธ๋จผํธ(88๊ฐ)๋ฅผ ์ฌ์ฉํ์ฌ 3๊ฐ emotion2vec variant์
+latency, RAM, ์์ธก ํ์ง์ ์ค์ธก ๋น๊ตํฉ๋๋ค.
+
+Usage:
+ python scripts/benchmark_emotion2vec.py # 3๊ฐ ์ ๋ถ
+ python scripts/benchmark_emotion2vec.py --variants plus_base # ๋จ์ผ
+ python scripts/benchmark_emotion2vec.py --device cuda # GPU
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import glob
+import json
+import logging
+import os
+import statistics
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import psutil
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Constants
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+VARIANT_CONFIGS = {
+ "base": "iic/emotion2vec_base",
+ "plus_base": "iic/emotion2vec_plus_base",
+ "plus_large": "iic/emotion2vec_plus_large",
+}
+
+# emotion2vec 9-class โ project 7-class mapping
+LABEL_MAP = {
+ "angry": "anger",
+ "disgusted": "disgust",
+ "fearful": "fear",
+ "happy": "joy",
+ "neutral": "neutral",
+ "sad": "sadness",
+ "surprised": "surprise",
+ "other": "neutral",
+ "unknown": "neutral",
+}
+
+PROJECT_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear", "disgust"]
+
+# Representative segments for Korean sanity check (indices into sorted segment list)
+# Will be selected dynamically: shortest, longest, and 3 evenly spaced
+SANITY_CHECK_COUNT = 5
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Segment Discovery
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def discover_segments(segments_dir: str) -> list[dict]:
+ """Find all segment WAV files and load metadata from stage1_output.json."""
+ pattern = os.path.join(segments_dir, "call_*", "seg_*.wav")
+ paths = sorted(glob.glob(pattern))
+
+ if not paths:
+ logger.error("No segment files found in %s", segments_dir)
+ sys.exit(1)
+
+ # Try to load metadata from stage1_output.json for text context
+ metadata = {}
+ stage1_path = Path(segments_dir).parent / "stage1_output.json"
+ if stage1_path.exists():
+ with open(stage1_path) as f:
+ data = json.load(f)
+ for seg in data.get("segments", []):
+ metadata[seg["audio_path"]] = {
+ "text": seg.get("text", ""),
+ "speaker_id": seg.get("speaker_id", ""),
+ "start": seg.get("start", 0),
+ "end": seg.get("end", 0),
+ }
+
+ segments = []
+ for p in paths:
+ call_id = Path(p).parent.name
+ seg_name = Path(p).stem
+ meta = metadata.get(p, {})
+ segments.append({
+ "path": p,
+ "call_id": call_id,
+ "seg_name": seg_name,
+ "text": meta.get("text", ""),
+ "speaker_id": meta.get("speaker_id", ""),
+ "duration_sec": meta.get("end", 0) - meta.get("start", 0),
+ })
+
+ logger.info("Discovered %d segments across %d calls",
+ len(segments), len(set(s["call_id"] for s in segments)))
+ return segments
+
+
+def select_sanity_segments(segments: list[dict], count: int = SANITY_CHECK_COUNT) -> list[dict]:
+ """Select representative segments for sanity check: shortest, longest, + evenly spaced."""
+ if len(segments) <= count:
+ return segments
+
+ # Sort by duration for selection
+ by_dur = sorted(segments, key=lambda s: s["duration_sec"])
+ # Filter to segments that have text (from stage1_output.json call)
+ with_text = [s for s in by_dur if s["text"]]
+ if len(with_text) < count:
+ with_text = by_dur
+
+ selected = [with_text[0], with_text[-1]] # shortest, longest
+ remaining = count - 2
+ step = max(1, len(with_text) // (remaining + 1))
+ for i in range(1, remaining + 1):
+ idx = min(i * step, len(with_text) - 1)
+ candidate = with_text[idx]
+ if candidate not in selected:
+ selected.append(candidate)
+
+ return selected[:count]
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Benchmarking
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def get_process_rss_mb() -> float:
+ """Current process RSS in MB."""
+ return psutil.Process(os.getpid()).memory_info().rss / (1024 * 1024)
+
+
+def map_predictions(raw_scores: dict[str, float]) -> dict:
+ """Map emotion2vec native labels to project 7-class taxonomy."""
+ mapped = {label: 0.0 for label in PROJECT_LABELS}
+ for native_label, score in raw_scores.items():
+ project_label = LABEL_MAP.get(native_label, "neutral")
+ mapped[project_label] += score
+
+ top_label = max(mapped, key=mapped.get)
+ return {
+ "label": top_label,
+ "confidence": mapped[top_label],
+ "scores": mapped,
+ "raw_scores": raw_scores,
+ }
+
+
+def benchmark_variant(
+ variant_name: str,
+ model_id: str,
+ segments: list[dict],
+ device: str = "cpu",
+ warmup: int = 3,
+) -> dict:
+ """Run full benchmark for one emotion2vec variant."""
+ from funasr import AutoModel
+
+ logger.info("=" * 60)
+ logger.info("Benchmarking: %s (%s)", variant_name, model_id)
+ logger.info("=" * 60)
+
+ result = {
+ "variant": variant_name,
+ "model_id": model_id,
+ "device": device,
+ }
+
+ # 1. Baseline RAM
+ gc.collect()
+ baseline_rss = get_process_rss_mb()
+
+ # 2. Model load + load time
+ logger.info("Loading model...")
+ load_start = time.perf_counter()
+ try:
+ model = AutoModel(model=model_id, device=device)
+ except Exception as e:
+ logger.error("Failed to load %s: %s", model_id, e)
+ result["error"] = str(e)
+ return result
+ load_time = time.perf_counter() - load_start
+ result["load_time_sec"] = round(load_time, 2)
+ logger.info("Model loaded in %.2fs", load_time)
+
+ # 3. Peak RAM after load
+ post_load_rss = get_process_rss_mb()
+ result["model_ram_mb"] = round(post_load_rss - baseline_rss, 1)
+ logger.info("Model RAM: %.1f MB", result["model_ram_mb"])
+
+ # 4. Warmup
+ logger.info("Warmup (%d runs)...", warmup)
+ warmup_segs = segments[:warmup] if len(segments) >= warmup else segments
+ for seg in warmup_segs:
+ try:
+ model.generate(seg["path"], granularity="utterance", extract_embedding=False)
+ except Exception as e:
+ logger.warning("Warmup failed on %s: %s", seg["path"], e)
+
+ # 5. Timed inference on all segments
+ logger.info("Running inference on %d segments...", len(segments))
+ predictions = []
+ latencies = []
+ errors = []
+ peak_rss = post_load_rss
+
+ for i, seg in enumerate(segments):
+ try:
+ t0 = time.perf_counter()
+ output = model.generate(
+ seg["path"], granularity="utterance", extract_embedding=False,
+ )
+ t1 = time.perf_counter()
+
+ latency_ms = (t1 - t0) * 1000
+ latencies.append(latency_ms)
+
+ # Parse emotion2vec output
+ raw_scores = {}
+ if output and isinstance(output, list) and len(output) > 0:
+ rec = output[0]
+ labels = rec.get("labels", [])
+ scores = rec.get("scores", [])
+ for label, score in zip(labels, scores):
+ raw_scores[label] = float(score)
+
+ mapped = map_predictions(raw_scores)
+
+ predictions.append({
+ "seg_name": seg["seg_name"],
+ "call_id": seg["call_id"],
+ "text": seg["text"],
+ "speaker_id": seg["speaker_id"],
+ "duration_sec": seg["duration_sec"],
+ "latency_ms": round(latency_ms, 1),
+ **mapped,
+ })
+
+ except Exception as e:
+ errors.append({"seg_name": seg["seg_name"], "error": str(e)})
+ logger.warning("Inference error on %s: %s", seg["seg_name"], e)
+
+ # Track peak RAM
+ current_rss = get_process_rss_mb()
+ peak_rss = max(peak_rss, current_rss)
+
+ if (i + 1) % 20 == 0:
+ logger.info(" %d/%d segments done", i + 1, len(segments))
+
+ # 6. Aggregate results
+ result["peak_ram_mb"] = round(peak_rss - baseline_rss, 1)
+ result["total_segments"] = len(segments)
+ result["successful"] = len(predictions)
+ result["errors"] = errors
+
+ if latencies:
+ result["latency"] = {
+ "mean_ms": round(statistics.mean(latencies), 1),
+ "median_ms": round(statistics.median(latencies), 1),
+ "std_ms": round(statistics.stdev(latencies), 1) if len(latencies) > 1 else 0,
+ "p95_ms": round(sorted(latencies)[int(len(latencies) * 0.95)], 1),
+ "min_ms": round(min(latencies), 1),
+ "max_ms": round(max(latencies), 1),
+ }
+ else:
+ result["latency"] = {}
+
+ # Emotion distribution
+ dist = {label: 0 for label in PROJECT_LABELS}
+ for pred in predictions:
+ dist[pred["label"]] += 1
+ result["emotion_distribution"] = dist
+
+ result["predictions"] = predictions
+
+ # 7. Cleanup
+ logger.info("Cleaning up model...")
+ del model
+ gc.collect()
+ try:
+ import torch
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+ except ImportError:
+ pass
+
+ logger.info("Done: %s โ mean latency %.1fms, peak RAM %.1fMB",
+ variant_name,
+ result.get("latency", {}).get("mean_ms", 0),
+ result.get("peak_ram_mb", 0))
+
+ return result
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Output Formatting
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def fmt_table(headers: list[str], rows: list[list[str]], col_widths: list[int] | None = None) -> str:
+ """Simple table formatter."""
+ if not col_widths:
+ col_widths = []
+ for i, h in enumerate(headers):
+ max_w = len(h)
+ for row in rows:
+ if i < len(row):
+ max_w = max(max_w, len(str(row[i])))
+ col_widths.append(max_w + 2)
+
+ def fmt_row(cells):
+ return "โ " + " โ ".join(str(c).ljust(w) for c, w in zip(cells, col_widths)) + " โ"
+
+ separator = "โโ" + "โโผโ".join("โ" * w for w in col_widths) + "โโค"
+ top = "โโ" + "โโฌโ".join("โ" * w for w in col_widths) + "โโ"
+ bottom = "โโ" + "โโดโ".join("โ" * w for w in col_widths) + "โโ"
+
+ lines = [top, fmt_row(headers), separator]
+ for row in rows:
+ lines.append(fmt_row(row))
+ lines.append(bottom)
+ return "\n".join(lines)
+
+
+def format_results(all_results: dict[str, dict], segments: list[dict]) -> str:
+ """Format benchmark results into readable console output."""
+ output_parts = []
+
+ # โโ Performance Comparison โโ
+ output_parts.append("\n=== Performance Comparison ===")
+ headers = ["Variant", "Load (s)", "RAM (MB)", "Latency mean (ms)", "Latency p95 (ms)", "Errors"]
+ rows = []
+ for name, res in all_results.items():
+ if "error" in res:
+ rows.append([name, "FAIL", "-", "-", "-", res["error"][:40]])
+ continue
+ lat = res.get("latency", {})
+ mean_str = f"{lat.get('mean_ms', 0):.1f} ยฑ {lat.get('std_ms', 0):.1f}"
+ rows.append([
+ name,
+ f"{res.get('load_time_sec', 0):.1f}",
+ f"{res.get('peak_ram_mb', 0):.0f}",
+ mean_str,
+ f"{lat.get('p95_ms', 0):.1f}",
+ str(len(res.get("errors", []))),
+ ])
+ output_parts.append(fmt_table(headers, rows))
+
+ # โโ Knockout Check โโ
+ output_parts.append("\n=== Knockout Check ===")
+ for name, res in all_results.items():
+ if "error" in res:
+ output_parts.append(f" {name}: โ LOAD FAILED")
+ continue
+ lat_mean = res.get("latency", {}).get("mean_ms", 999)
+ ram = res.get("peak_ram_mb", 999)
+ lat_ok = "โ
" if lat_mean <= 500 else "โ"
+ ram_ok = "โ
" if ram <= 2048 else "โ"
+ output_parts.append(f" {name}: Latency {lat_ok} ({lat_mean:.0f}ms โค 500ms) RAM {ram_ok} ({ram:.0f}MB โค 2048MB)")
+
+ # โโ Emotion Distribution โโ
+ output_parts.append("\n=== Emotion Distribution (across all segments) ===")
+ headers = ["Variant"] + PROJECT_LABELS
+ rows = []
+ for name, res in all_results.items():
+ if "error" in res:
+ continue
+ dist = res.get("emotion_distribution", {})
+ rows.append([name] + [str(dist.get(l, 0)) for l in PROJECT_LABELS])
+ output_parts.append(fmt_table(headers, rows))
+
+ # โโ Korean Sanity Check โโ
+ output_parts.append("\n=== Korean Sanity Check ===")
+ sanity_segs = select_sanity_segments(segments)
+ for seg in sanity_segs:
+ text_preview = seg["text"][:50] + "..." if len(seg["text"]) > 50 else seg["text"]
+ output_parts.append(f'\n {seg["seg_name"]} ({seg["duration_sec"]:.1f}s): "{text_preview}"')
+ for name, res in all_results.items():
+ if "error" in res:
+ output_parts.append(f" {name}: FAILED")
+ continue
+ # Find matching prediction
+ preds = res.get("predictions", [])
+ match = next((p for p in preds if p["seg_name"] == seg["seg_name"]), None)
+ if match:
+ output_parts.append(f" {name:12s}: {match['label']:10s} ({match['confidence']:.2f})")
+ else:
+ output_parts.append(f" {name:12s}: no prediction")
+
+ # โโ Variant Agreement โโ
+ output_parts.append("\n=== Variant Agreement ===")
+ valid_results = {k: v for k, v in all_results.items() if "error" not in v}
+ if len(valid_results) >= 2:
+ variant_names = list(valid_results.keys())
+ # Build prediction maps: seg_name -> label
+ pred_maps = {}
+ for name, res in valid_results.items():
+ pred_maps[name] = {p["seg_name"]: p["label"] for p in res.get("predictions", [])}
+
+ # All-agree count
+ all_seg_names = set()
+ for pm in pred_maps.values():
+ all_seg_names.update(pm.keys())
+
+ agree_count = 0
+ total_count = 0
+ for seg_name in all_seg_names:
+ labels = [pm.get(seg_name) for pm in pred_maps.values() if seg_name in pm]
+ if len(labels) == len(valid_results):
+ total_count += 1
+ if len(set(labels)) == 1:
+ agree_count += 1
+
+ output_parts.append(f" All {len(valid_results)} variants agree: {agree_count}/{total_count} ({agree_count/max(total_count,1)*100:.0f}%)")
+
+ # Pairwise agreement
+ for i in range(len(variant_names)):
+ for j in range(i + 1, len(variant_names)):
+ a, b = variant_names[i], variant_names[j]
+ common = set(pred_maps[a].keys()) & set(pred_maps[b].keys())
+ pair_agree = sum(1 for s in common if pred_maps[a][s] == pred_maps[b][s])
+ pct = pair_agree / max(len(common), 1) * 100
+ output_parts.append(f" {a} vs {b}: {pair_agree}/{len(common)} ({pct:.0f}%)")
+
+ return "\n".join(output_parts)
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Main
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def main():
+ parser = argparse.ArgumentParser(description="emotion2vec variant benchmark")
+ parser.add_argument(
+ "--variants", nargs="*", default=list(VARIANT_CONFIGS.keys()),
+ choices=list(VARIANT_CONFIGS.keys()),
+ help="Which variants to benchmark (default: all)",
+ )
+ parser.add_argument("--segments-dir", default="data/segments", help="Segments directory")
+ parser.add_argument("--output-json", default="data/benchmark_results.json", help="Output JSON path")
+ parser.add_argument("--device", default="cpu", choices=["cpu", "cuda"], help="Compute device")
+ parser.add_argument("--warmup", type=int, default=3, help="Warmup iterations")
+ args = parser.parse_args()
+
+ # Check dependency
+ try:
+ import funasr # noqa: F401
+ except ImportError:
+ logger.error("funasr not installed. Run: pip install funasr onnxruntime")
+ sys.exit(1)
+
+ # Discover segments
+ segments = discover_segments(args.segments_dir)
+ logger.info("Total segments: %d", len(segments))
+
+ # Run benchmarks
+ all_results = {}
+ for variant_name in args.variants:
+ model_id = VARIANT_CONFIGS[variant_name]
+ result = benchmark_variant(
+ variant_name, model_id, segments,
+ device=args.device, warmup=args.warmup,
+ )
+ all_results[variant_name] = result
+
+ # Format and print results
+ report = format_results(all_results, segments)
+ print(report)
+
+ # Save JSON
+ output_path = Path(args.output_json)
+ output_path.parent.mkdir(parents=True, exist_ok=True)
+
+ # Get system info
+ import platform
+ try:
+ import torch
+ torch_version = torch.__version__
+ cuda_available = torch.cuda.is_available()
+ except ImportError:
+ torch_version = "not installed"
+ cuda_available = False
+
+ output_data = {
+ "metadata": {
+ "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
+ "device": args.device,
+ "total_segments": len(segments),
+ "python_version": platform.python_version(),
+ "torch_version": torch_version,
+ "cuda_available": cuda_available,
+ "cpu": platform.processor() or "unknown",
+ "ram_total_gb": round(psutil.virtual_memory().total / (1024**3), 1),
+ },
+ "results": all_results,
+ }
+
+ with open(output_path, "w", encoding="utf-8") as f:
+ json.dump(output_data, f, indent=2, ensure_ascii=False, default=str)
+
+ logger.info("Results saved to %s", output_path)
+ print(f"\nFull results saved to {output_path}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/benchmark_ser_models.py b/scripts/benchmark_ser_models.py
new file mode 100644
index 0000000000000000000000000000000000000000..b1a32225f5d7da87eea0dd7bbe03d991d1ee97d8
--- /dev/null
+++ b/scripts/benchmark_ser_models.py
@@ -0,0 +1,799 @@
+#!/usr/bin/env python3
+"""3-Model SER Benchmark โ emotion2vec vs SpeechBrain vs Whisper+Head.
+
+AI Hub ํ๊ตญ์ด ๊ฐ์ ๋ฐ์ดํฐ์
ํ
์คํธ ์๋ธ์
์ ์ฌ์ฉํ์ฌ 3๊ฐ ๋ชจ๋ธ์
+์ ํ๋, ๋ ์ดํด์, ๋ฉ๋ชจ๋ฆฌ ์ฌ์ฉ๋์ ๊ฐ๊ด์ ์ผ๋ก ๋น๊ตํ๋ค.
+
+Usage:
+ # 2๊ฐ ๋ชจ๋ธ ๋จผ์ (Whisper head ์์ด)
+ python scripts/benchmark_ser_models.py \\
+ --test-dir data/evaluation/korean \\
+ --models emotion2vec speechbrain
+
+ # ์ ์ฒด 3๊ฐ ๋ชจ๋ธ
+ python scripts/benchmark_ser_models.py \\
+ --test-dir data/evaluation/korean \\
+ --models emotion2vec speechbrain whisper \\
+ --whisper-head-ckpt data/models/whisper_emotion_head.pt
+
+ # Quick smoke test
+ python scripts/benchmark_ser_models.py \\
+ --test-dir data/evaluation/korean \\
+ --models emotion2vec --max-samples 10
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import gc
+import json
+import logging
+import os
+import statistics
+import sys
+import tempfile
+import time
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+import numpy as np
+import psutil
+import soundfile as sf
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Constants
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+EVAL_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear"]
+
+# Knockout criteria (from evaluation-framework.md)
+KNOCKOUT_F1 = 0.70
+KNOCKOUT_LATENCY_MS = 500
+KNOCKOUT_RAM_MB = 2048
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Model Adapter Interface
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+class SERModelAdapter(ABC):
+ """Abstract base for SER model adapters."""
+
+ name: str
+ model_id: str
+ params_m: int # millions
+
+ @abstractmethod
+ def load(self, device: str) -> None:
+ ...
+
+ @abstractmethod
+ def predict(self, audio_path: str) -> dict[str, float]:
+ """Return {emotion_label: score} in project taxonomy."""
+ ...
+
+ @abstractmethod
+ def unload(self) -> None:
+ ...
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Adapter 1: emotion2vec_plus_base
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+class Emotion2vecAdapter(SERModelAdapter):
+ name = "emotion2vec_plus_base"
+ model_id = "iic/emotion2vec_plus_base"
+ params_m = 90
+
+ # emotion2vec 9-class โ project 7-class (from src/stage2/audio_emotion.py)
+ LABEL_MAP = {
+ "angry": "anger", "disgusted": "disgust", "fearful": "fear",
+ "happy": "joy", "neutral": "neutral", "sad": "sadness",
+ "surprised": "surprise", "other": "neutral", "unknown": "neutral",
+ "็ๆฐ/angry": "anger", "ๅๆถ/disgusted": "disgust",
+ "ๆๆง/fearful": "fear", "ๅผๅฟ/happy": "joy",
+ "ไธญ็ซ/neutral": "neutral", "้พ่ฟ/sad": "sadness",
+ "ๅๆ/surprised": "surprise", "ๅ
ถไป/other": "neutral", "": "neutral",
+ }
+
+ def __init__(self):
+ self._model = None
+
+ def load(self, device: str) -> None:
+ from funasr import AutoModel
+ self._model = AutoModel(model=self.model_id, device=device, hub="hf")
+
+ def predict(self, audio_path: str) -> dict[str, float]:
+ output = self._model.generate(
+ audio_path, granularity="utterance", extract_embedding=False,
+ )
+ scores = {label: 0.0 for label in EVAL_LABELS}
+ if output and isinstance(output, list) and len(output) > 0:
+ rec = output[0]
+ for native_label, score in zip(rec.get("labels", []), rec.get("scores", [])):
+ mapped = self.LABEL_MAP.get(native_label, "neutral")
+ if mapped in scores:
+ scores[mapped] += float(score)
+ # Normalize
+ total = sum(scores.values())
+ if total > 0:
+ scores = {k: v / total for k, v in scores.items()}
+ return scores
+
+ def unload(self) -> None:
+ del self._model
+ self._model = None
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Adapter 2: SpeechBrain wav2vec2-IEMOCAP
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+class SpeechBrainAdapter(SERModelAdapter):
+ name = "speechbrain_wav2vec2"
+ model_id = "speechbrain/emotion-recognition-wav2vec2-IEMOCAP"
+ params_m = 314
+
+ # SpeechBrain 4-class โ project taxonomy
+ # NOTE: This model CANNOT predict fear or surprise
+ LABEL_MAP = {
+ "ang": "anger",
+ "hap": "joy",
+ "sad": "sadness",
+ "neu": "neutral",
+ }
+
+ def __init__(self):
+ self._classifier = None
+ self._label_order = None # populated from label_encoder
+
+ def load(self, device: str) -> None:
+ import torch
+ from speechbrain.inference.classifiers import EncoderClassifier
+ self._classifier = EncoderClassifier.from_hparams(
+ source=self.model_id,
+ run_opts={"device": device},
+ )
+ self._classifier = self._classifier.to(device)
+
+ # Get label order from label_encoder
+ try:
+ le = self._classifier.hparams.label_encoder
+ # lab2ind: {'neu': 0, 'ang': 1, 'hap': 2, 'sad': 3}
+ self._label_order = [None] * len(le.lab2ind)
+ for lab, idx in le.lab2ind.items():
+ self._label_order[idx] = lab
+ logger.info("SpeechBrain labels: %s", self._label_order)
+ except Exception:
+ self._label_order = ["neu", "ang", "hap", "sad"]
+
+ def predict(self, audio_path: str) -> dict[str, float]:
+ import torch
+ import torchaudio
+
+ signal, sr = torchaudio.load(audio_path)
+ if sr != 16000:
+ signal = torchaudio.functional.resample(signal, sr, 16000)
+ if signal.shape[0] > 1:
+ signal = signal.mean(dim=0, keepdim=True)
+
+ # Use modules directly (classify_batch broken in SpeechBrain 1.1.0)
+ with torch.no_grad():
+ feats = self._classifier.mods.wav2vec2(signal)
+ pooled = self._classifier.mods.avg_pool(feats)
+ logits = self._classifier.mods.output_mlp(pooled)
+ probs = torch.softmax(logits.squeeze(1), dim=-1).squeeze().tolist()
+
+ if isinstance(probs, float):
+ probs = [probs]
+
+ scores = {label: 0.0 for label in EVAL_LABELS}
+ for sb_label, prob in zip(self._label_order, probs):
+ mapped = self.LABEL_MAP.get(sb_label, "neutral")
+ if mapped in scores:
+ scores[mapped] += prob
+
+ return scores
+
+ def unload(self) -> None:
+ del self._classifier
+ self._classifier = None
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Adapter 3: Whisper-Medium + Emotion Head
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+class WhisperMediumAdapter(SERModelAdapter):
+ name = "whisper_medium_head"
+ model_id = "openai/whisper-medium"
+ params_m = 769
+
+ def __init__(self, head_ckpt: str | None = None):
+ self._encoder = None
+ self._head = None
+ self._processor = None
+ self._head_ckpt = head_ckpt
+ self._device = "cpu"
+
+ def load(self, device: str) -> None:
+ import torch
+ from transformers import WhisperModel, WhisperFeatureExtractor
+
+ self._device = device
+ self._processor = WhisperFeatureExtractor.from_pretrained(self.model_id)
+ self._encoder = WhisperModel.from_pretrained(self.model_id).to(device)
+ self._encoder.eval()
+
+ # Classifier head: hidden_dim โ 6 classes
+ hidden_dim = self._encoder.config.d_model # 1024 for medium
+ self._head = torch.nn.Linear(hidden_dim, len(EVAL_LABELS)).to(device)
+
+ if self._head_ckpt and Path(self._head_ckpt).exists():
+ logger.info("Loading Whisper emotion head from %s", self._head_ckpt)
+ state = torch.load(self._head_ckpt, map_location=device, weights_only=True)
+ self._head.load_state_dict(state)
+ else:
+ logger.warning("No trained Whisper head โ using random weights (baseline)")
+
+ self._head.eval()
+
+ def predict(self, audio_path: str) -> dict[str, float]:
+ import torch
+ import librosa
+
+ # Load and preprocess
+ audio, sr = librosa.load(audio_path, sr=16000)
+ inputs = self._processor(
+ audio, sampling_rate=16000, return_tensors="pt",
+ )
+ input_features = inputs.input_features.to(self._device)
+
+ with torch.no_grad():
+ encoder_out = self._encoder.encoder(input_features)
+ hidden = encoder_out.last_hidden_state # (1, T, D)
+ pooled = hidden.mean(dim=1) # (1, D)
+ logits = self._head(pooled) # (1, 6)
+ probs = torch.softmax(logits, dim=-1).squeeze().cpu().tolist()
+
+ scores = {}
+ for label, prob in zip(EVAL_LABELS, probs):
+ scores[label] = prob
+ return scores
+
+ def unload(self) -> None:
+ del self._encoder, self._head, self._processor
+ self._encoder = self._head = self._processor = None
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Phone Augmentation
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def apply_phone_augmentation(audio_path: str) -> str:
+ """Apply phone-quality degradation, return path to temp WAV file."""
+ sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+ from common.phone_simulator import PhoneSimulator, CompandingType
+
+ audio, sr = sf.read(audio_path, dtype="float32")
+ if audio.ndim == 2:
+ audio = audio.mean(axis=1)
+
+ sim = PhoneSimulator(companding=CompandingType.ALAW)
+ degraded, new_sr = sim.process(audio, sr)
+
+ # Save to temp file
+ tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+ sf.write(tmp.name, degraded, new_sr, subtype="PCM_16")
+ return tmp.name
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Test Data Loading
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def load_test_data(test_dir: str, max_samples: int | None = None) -> list[dict]:
+ """Load test samples from prepared subset."""
+ csv_path = Path(test_dir) / "test_labels.csv"
+ if not csv_path.exists():
+ logger.error("test_labels.csv not found in %s", test_dir)
+ sys.exit(1)
+
+ samples = []
+ with open(csv_path, encoding="utf-8") as f:
+ reader = csv.DictReader(f)
+ for row in reader:
+ audio_path = str(Path(test_dir) / row["file_path"])
+ if not Path(audio_path).exists():
+ logger.warning("Audio file not found: %s", audio_path)
+ continue
+ samples.append({
+ "audio_path": audio_path,
+ "emotion": row["emotion"],
+ "duration": float(row["duration"]),
+ "speaker_id": row.get("speaker_id", ""),
+ "intensity": row.get("intensity", ""),
+ })
+
+ if max_samples and len(samples) > max_samples:
+ import random
+ random.seed(42)
+ samples = random.sample(samples, max_samples)
+
+ logger.info("Loaded %d test samples from %s", len(samples), test_dir)
+ return samples
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Benchmark Runner
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def get_process_rss_mb() -> float:
+ return psutil.Process(os.getpid()).memory_info().rss / (1024 * 1024)
+
+
+def benchmark_model(
+ adapter: SERModelAdapter,
+ samples: list[dict],
+ device: str,
+ phone_augment: bool,
+ warmup: int = 5,
+) -> dict:
+ """Run full benchmark for one model on both clean and optionally phone conditions."""
+ logger.info("=" * 60)
+ logger.info("Benchmarking: %s (%dM params)", adapter.name, adapter.params_m)
+ logger.info("=" * 60)
+
+ result = {
+ "model": adapter.name,
+ "model_id": adapter.model_id,
+ "params_m": adapter.params_m,
+ "device": device,
+ }
+
+ # Baseline RAM
+ gc.collect()
+ baseline_rss = get_process_rss_mb()
+
+ # Load model
+ logger.info("Loading model...")
+ load_start = time.perf_counter()
+ try:
+ adapter.load(device)
+ except Exception as e:
+ logger.error("Failed to load %s: %s", adapter.name, e)
+ result["error"] = str(e)
+ return result
+ load_time = time.perf_counter() - load_start
+ result["load_time_sec"] = round(load_time, 2)
+
+ post_load_rss = get_process_rss_mb()
+ result["model_ram_mb"] = round(post_load_rss - baseline_rss, 1)
+ logger.info("Loaded in %.1fs, RAM: %.0fMB", load_time, result["model_ram_mb"])
+
+ # Run for each condition
+ conditions = ["clean"]
+ if phone_augment:
+ conditions.append("phone")
+
+ for condition in conditions:
+ logger.info("--- Condition: %s ---", condition)
+
+ # Warmup
+ warmup_samples = samples[:warmup] if len(samples) >= warmup else samples
+ for s in warmup_samples:
+ try:
+ audio_path = s["audio_path"]
+ if condition == "phone":
+ audio_path = apply_phone_augmentation(audio_path)
+ adapter.predict(audio_path)
+ if condition == "phone":
+ os.unlink(audio_path)
+ except Exception:
+ pass
+
+ # Inference
+ y_true = []
+ y_pred = []
+ latencies = []
+ errors = []
+ peak_rss = get_process_rss_mb()
+
+ for i, sample in enumerate(samples):
+ audio_path = sample["audio_path"]
+ tmp_path = None
+
+ try:
+ if condition == "phone":
+ tmp_path = apply_phone_augmentation(audio_path)
+ audio_path = tmp_path
+
+ t0 = time.perf_counter()
+ scores = adapter.predict(audio_path)
+ t1 = time.perf_counter()
+
+ latency_ms = (t1 - t0) * 1000
+ latencies.append(latency_ms)
+
+ pred_label = max(scores, key=scores.get)
+ y_true.append(sample["emotion"])
+ y_pred.append(pred_label)
+
+ except Exception as e:
+ errors.append({"index": i, "error": str(e)})
+ logger.warning("Error on sample %d: %s", i, e)
+ finally:
+ if tmp_path and os.path.exists(tmp_path):
+ os.unlink(tmp_path)
+
+ current_rss = get_process_rss_mb()
+ peak_rss = max(peak_rss, current_rss)
+
+ if (i + 1) % 50 == 0:
+ logger.info(" %d/%d done (mean lat: %.0fms)", i + 1, len(samples),
+ statistics.mean(latencies) if latencies else 0)
+
+ # Compute metrics
+ cond_result = compute_metrics(y_true, y_pred, latencies, peak_rss - baseline_rss, errors)
+ result[condition] = cond_result
+
+ logger.info(" %s: macro_f1=%.3f, accuracy=%.3f, mean_latency=%.0fms, peak_ram=%.0fMB",
+ condition,
+ cond_result["macro_f1"],
+ cond_result["accuracy"],
+ cond_result["latency"]["mean_ms"],
+ cond_result["peak_ram_mb"])
+
+ # Unload
+ adapter.unload()
+ gc.collect()
+ try:
+ import torch
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+ except ImportError:
+ pass
+
+ return result
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Metrics
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def compute_metrics(
+ y_true: list[str],
+ y_pred: list[str],
+ latencies: list[float],
+ peak_ram_mb: float,
+ errors: list[dict],
+) -> dict:
+ """Compute accuracy, F1, confusion matrix, latency stats."""
+ from sklearn.metrics import (
+ accuracy_score,
+ precision_recall_fscore_support,
+ confusion_matrix,
+ )
+
+ if not y_true or not y_pred:
+ return {
+ "accuracy": 0.0, "macro_f1": 0.0, "weighted_f1": 0.0,
+ "per_class": {l: {"precision": 0, "recall": 0, "f1": 0, "support": 0} for l in EVAL_LABELS},
+ "confusion_matrix": [[0] * len(EVAL_LABELS)] * len(EVAL_LABELS),
+ "confusion_labels": EVAL_LABELS,
+ "latency": {}, "peak_ram_mb": round(peak_ram_mb, 1),
+ "total_samples": 0, "errors": errors,
+ "note": "All samples failed โ no predictions available",
+ }
+
+ accuracy = accuracy_score(y_true, y_pred)
+ precision, recall, f1, support = precision_recall_fscore_support(
+ y_true, y_pred, labels=EVAL_LABELS, average=None, zero_division=0,
+ )
+ macro_f1 = float(np.mean(f1))
+ weighted_f1 = float(np.average(f1, weights=support)) if sum(support) > 0 else 0.0
+
+ cm = confusion_matrix(y_true, y_pred, labels=EVAL_LABELS).tolist()
+
+ per_class = {}
+ for i, label in enumerate(EVAL_LABELS):
+ per_class[label] = {
+ "precision": round(float(precision[i]), 4),
+ "recall": round(float(recall[i]), 4),
+ "f1": round(float(f1[i]), 4),
+ "support": int(support[i]),
+ }
+
+ latency_stats = {}
+ if latencies:
+ latency_stats = {
+ "mean_ms": round(statistics.mean(latencies), 1),
+ "median_ms": round(statistics.median(latencies), 1),
+ "std_ms": round(statistics.stdev(latencies), 1) if len(latencies) > 1 else 0,
+ "p95_ms": round(sorted(latencies)[int(len(latencies) * 0.95)], 1),
+ "min_ms": round(min(latencies), 1),
+ "max_ms": round(max(latencies), 1),
+ }
+
+ return {
+ "accuracy": round(accuracy, 4),
+ "macro_f1": round(macro_f1, 4),
+ "weighted_f1": round(weighted_f1, 4),
+ "per_class": per_class,
+ "confusion_matrix": cm,
+ "confusion_labels": EVAL_LABELS,
+ "latency": latency_stats,
+ "peak_ram_mb": round(peak_ram_mb, 1),
+ "total_samples": len(y_true),
+ "errors": errors,
+ }
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Knockout Check
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def knockout_check(result: dict) -> dict:
+ """Check if model passes knockout criteria."""
+ checks = {}
+ for condition in ["clean", "phone"]:
+ if condition not in result:
+ continue
+ cond = result[condition]
+ f1_ok = cond["macro_f1"] >= KNOCKOUT_F1
+ lat_ok = cond["latency"].get("mean_ms", 999) <= KNOCKOUT_LATENCY_MS
+ ram_ok = cond["peak_ram_mb"] <= KNOCKOUT_RAM_MB
+ checks[condition] = {
+ "korean_f1": f"{'PASS' if f1_ok else 'FAIL'} ({cond['macro_f1']:.3f} {'โฅ' if f1_ok else '<'} {KNOCKOUT_F1})",
+ "latency": f"{'PASS' if lat_ok else 'FAIL'} ({cond['latency'].get('mean_ms', 0):.0f}ms {'โค' if lat_ok else '>'} {KNOCKOUT_LATENCY_MS}ms)",
+ "ram": f"{'PASS' if ram_ok else 'FAIL'} ({cond['peak_ram_mb']:.0f}MB {'โค' if ram_ok else '>'} {KNOCKOUT_RAM_MB}MB)",
+ "overall": "PASS" if (f1_ok and lat_ok and ram_ok) else "FAIL",
+ }
+ return checks
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Report Generation
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def generate_markdown_report(all_results: dict, output_path: str):
+ """Generate a markdown comparison report."""
+ lines = [
+ "# 3-Model SER Benchmark Report",
+ "",
+ f"**Generated**: {time.strftime('%Y-%m-%d %H:%M:%S')}",
+ f"**Dataset**: AI Hub #71631 (๊ฐ์ ์ด ํ๊น
๋ ์์ ๋ํ - ์ฑ์ธ)",
+ f"**Evaluation Classes**: {', '.join(EVAL_LABELS)} (6-class, no disgust)",
+ "",
+ "---",
+ "",
+ "## Summary Comparison",
+ "",
+ ]
+
+ # Summary table
+ headers = ["Model", "Params", "Clean F1", "Phone F1", "Latency (mean)", "Latency (p95)", "RAM", "Knockout"]
+ rows = []
+ for name, res in all_results.items():
+ if "error" in res:
+ rows.append(f"| {name} | {res.get('params_m', '?')}M | LOAD FAILED | - | - | - | - | FAIL |")
+ continue
+ clean = res.get("clean", {})
+ phone = res.get("phone", {})
+ ko = knockout_check(res)
+ clean_ko = ko.get("clean", {}).get("overall", "N/A")
+ rows.append(
+ f"| {name} | {res['params_m']}M "
+ f"| {clean.get('macro_f1', 0):.3f} "
+ f"| {phone.get('macro_f1', 'N/A') if phone else 'N/A'} "
+ f"| {clean.get('latency', {}).get('mean_ms', 0):.0f}ms "
+ f"| {clean.get('latency', {}).get('p95_ms', 0):.0f}ms "
+ f"| {clean.get('peak_ram_mb', 0):.0f}MB "
+ f"| {clean_ko} |"
+ )
+
+ lines.append(f"| {' | '.join(headers)} |")
+ lines.append(f"| {'---|' * len(headers)}")
+ lines.extend(rows)
+ lines.append("")
+
+ # Knockout details
+ lines.extend(["", "## Knockout Check", ""])
+ for name, res in all_results.items():
+ if "error" in res:
+ continue
+ ko = knockout_check(res)
+ lines.append(f"### {name}")
+ for cond, checks in ko.items():
+ lines.append(f"**{cond}**: {checks['overall']}")
+ lines.append(f" - F1: {checks['korean_f1']}")
+ lines.append(f" - Latency: {checks['latency']}")
+ lines.append(f" - RAM: {checks['ram']}")
+ lines.append("")
+
+ # Per-model details with confusion matrix
+ lines.extend(["## Per-Model Details", ""])
+ for name, res in all_results.items():
+ if "error" in res:
+ continue
+ lines.append(f"### {name}")
+
+ for condition in ["clean", "phone"]:
+ if condition not in res:
+ continue
+ cond = res[condition]
+ lines.extend([
+ f"",
+ f"#### {condition.title()} Condition",
+ f"",
+ f"- Accuracy: {cond['accuracy']:.3f}",
+ f"- Macro F1: {cond['macro_f1']:.3f}",
+ f"- Weighted F1: {cond['weighted_f1']:.3f}",
+ f"",
+ "**Per-class F1:**",
+ "",
+ "| Emotion | Precision | Recall | F1 | Support |",
+ "|---|---|---|---|---|",
+ ])
+ for label in EVAL_LABELS:
+ pc = cond["per_class"].get(label, {})
+ lines.append(
+ f"| {label} | {pc.get('precision', 0):.3f} "
+ f"| {pc.get('recall', 0):.3f} "
+ f"| {pc.get('f1', 0):.3f} "
+ f"| {pc.get('support', 0)} |"
+ )
+
+ # Confusion matrix
+ lines.extend(["", "**Confusion Matrix:**", ""])
+ cm = cond.get("confusion_matrix", [])
+ if cm:
+ lines.append("| | " + " | ".join(EVAL_LABELS) + " |")
+ lines.append("| --- | " + " | ".join(["---"] * len(EVAL_LABELS)) + " |")
+ for i, row in enumerate(cm):
+ lines.append(f"| **{EVAL_LABELS[i]}** | " + " | ".join(str(v) for v in row) + " |")
+
+ lines.append("")
+
+ # Limitations
+ lines.extend([
+ "## Known Limitations",
+ "",
+ "- **SpeechBrain wav2vec2-IEMOCAP**: Only outputs 4 classes (angry, happy, sad, neutral). "
+ "Cannot predict fear or surprise โ structurally penalized in 6-class macro F1.",
+ "- **Whisper-Medium + Head**: Requires a separately trained classifier head. "
+ "Without training, results reflect random baseline (~16.7%).",
+ "- **AI Hub dataset**: No 'disgust' class โ evaluated as 6-class instead of project's 7-class.",
+ "",
+ ])
+
+ Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+ with open(output_path, "w", encoding="utf-8") as f:
+ f.write("\n".join(lines))
+ logger.info("Markdown report saved to %s", output_path)
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Main
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+ADAPTER_MAP = {
+ "emotion2vec": Emotion2vecAdapter,
+ "speechbrain": SpeechBrainAdapter,
+ "whisper": WhisperMediumAdapter,
+}
+
+
+def main():
+ parser = argparse.ArgumentParser(description="3-Model SER Benchmark")
+ parser.add_argument("--test-dir", required=True, help="ํ
์คํธ ์๋ธ์
๋๋ ํ ๋ฆฌ (test_labels.csv ํฌํจ)")
+ parser.add_argument("--models", nargs="+", default=["emotion2vec", "speechbrain"],
+ choices=list(ADAPTER_MAP.keys()), help="๋ฒค์น๋งํฌํ ๋ชจ๋ธ (default: emotion2vec speechbrain)")
+ parser.add_argument("--device", default="cpu", choices=["cpu", "cuda"], help="Compute device")
+ parser.add_argument("--whisper-head-ckpt", default=None, help="Whisper emotion head ์ฒดํฌํฌ์ธํธ ๊ฒฝ๋ก")
+ parser.add_argument("--phone-augment", action="store_true", default=False, help="Phone augmentation ํ๊ฐ ์ถ๊ฐ")
+ parser.add_argument("--warmup", type=int, default=5, help="Warmup ํ์")
+ parser.add_argument("--max-samples", type=int, default=None, help="์ต๋ ์ํ ์ (smoke test์ฉ)")
+ parser.add_argument("--output-json", default="data/evaluation/benchmark_3model_results.json")
+ parser.add_argument("--output-md", default="docs/stage2/benchmark-3model-report.md")
+ args = parser.parse_args()
+
+ # Load test data
+ samples = load_test_data(args.test_dir, max_samples=args.max_samples)
+ if not samples:
+ logger.error("No test samples loaded")
+ sys.exit(1)
+
+ # Run benchmarks
+ all_results = {}
+ for model_name in args.models:
+ adapter_cls = ADAPTER_MAP[model_name]
+ if model_name == "whisper":
+ adapter = adapter_cls(head_ckpt=args.whisper_head_ckpt)
+ else:
+ adapter = adapter_cls()
+
+ result = benchmark_model(
+ adapter, samples, args.device,
+ phone_augment=args.phone_augment,
+ warmup=args.warmup,
+ )
+ result["knockout"] = knockout_check(result)
+ all_results[adapter.name] = result
+
+ # Save JSON
+ output_json_path = Path(args.output_json)
+ output_json_path.parent.mkdir(parents=True, exist_ok=True)
+
+ import platform
+ try:
+ import torch
+ torch_version = torch.__version__
+ cuda_available = torch.cuda.is_available()
+ except ImportError:
+ torch_version = "not installed"
+ cuda_available = False
+
+ output_data = {
+ "metadata": {
+ "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
+ "device": args.device,
+ "test_samples": len(samples),
+ "eval_classes": EVAL_LABELS,
+ "conditions": ["clean"] + (["phone"] if args.phone_augment else []),
+ "system_info": {
+ "cpu": platform.processor() or "unknown",
+ "ram_total_gb": round(psutil.virtual_memory().total / (1024**3), 1),
+ "python": platform.python_version(),
+ "torch": torch_version,
+ "cuda": cuda_available,
+ },
+ },
+ "results": all_results,
+ }
+
+ with open(output_json_path, "w", encoding="utf-8") as f:
+ json.dump(output_data, f, indent=2, ensure_ascii=False, default=str)
+ logger.info("JSON results saved to %s", output_json_path)
+
+ # Generate markdown report
+ generate_markdown_report(all_results, args.output_md)
+
+ # Console summary
+ print("\n" + "=" * 60)
+ print("BENCHMARK COMPLETE")
+ print("=" * 60)
+ for name, res in all_results.items():
+ if "error" in res:
+ print(f"\n {name}: LOAD FAILED โ {res['error']}")
+ continue
+ clean = res.get("clean", {})
+ ko = res.get("knockout", {}).get("clean", {})
+ print(f"\n {name} ({res['params_m']}M params):")
+ print(f" Clean F1: {clean.get('macro_f1', 0):.3f} Accuracy: {clean.get('accuracy', 0):.3f}")
+ print(f" Latency: {clean.get('latency', {}).get('mean_ms', 0):.0f}ms (mean), "
+ f"{clean.get('latency', {}).get('p95_ms', 0):.0f}ms (p95)")
+ print(f" RAM: {clean.get('peak_ram_mb', 0):.0f}MB")
+ print(f" Knockout: {ko.get('overall', 'N/A')}")
+
+ if args.phone_augment:
+ print("\n --- Phone Degradation ---")
+ for name, res in all_results.items():
+ if "error" in res or "phone" not in res:
+ continue
+ clean_f1 = res.get("clean", {}).get("macro_f1", 0)
+ phone_f1 = res["phone"]["macro_f1"]
+ drop = clean_f1 - phone_f1
+ print(f" {name}: {clean_f1:.3f} โ {phone_f1:.3f} (ฮ={drop:+.3f})")
+
+ print(f"\n Results: {args.output_json}")
+ print(f" Report: {args.output_md}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/build_english_fusion_manifest.py b/scripts/build_english_fusion_manifest.py
new file mode 100644
index 0000000000000000000000000000000000000000..97cd3c8807d4c4494ca169802496b8945bfd0c04
--- /dev/null
+++ b/scripts/build_english_fusion_manifest.py
@@ -0,0 +1,135 @@
+#!/usr/bin/env python3
+"""Build unified English fusion evaluation manifest.
+
+Combines:
+ - JL-Corpus 5-class (angryโanger, happyโjoy, sadโsadness, neutral, anxiousโfear)
+ - SAVEE 2-class (disgust, surprise) with WhisperX ASR transcripts
+ - MELD fusion disgust + surprise samples (natural dialogue)
+
+Output: data/english_fusion/manifest.json
+Format: [{"path": ..., "text": ..., "label": ..., "source": ..., "speaker": ...}, ...]
+"""
+from __future__ import annotations
+
+import json
+import logging
+import re
+from collections import Counter
+from pathlib import Path
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+PROJECT_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear", "disgust"]
+
+JL_MAP = {
+ "angry": "anger",
+ "happy": "joy",
+ "sad": "sadness",
+ "neutral": "neutral",
+ "anxious": "fear",
+}
+
+SAVEE_MAP = {"d": "disgust", "su": "surprise"}
+SAVEE_PATTERN = re.compile(r"^(?P[A-Z]{2})_(?Pd|su)(?P\d+)\.wav$")
+JL_PATTERN = re.compile(r"^(?P[a-z]+\d+)_(?P[a-z]+)_(?P[^_]+)_(?P\d+)\.wav$")
+
+OUT_DIR = Path("data/english_fusion")
+OUT_MANIFEST = OUT_DIR / "manifest.json"
+
+
+def load_jl_corpus() -> list[dict]:
+ jl_dir = Path("data/jl_corpus")
+ rows: list[dict] = []
+ for wav in sorted(jl_dir.glob("*.wav")):
+ m = JL_PATTERN.match(wav.name)
+ if not m:
+ continue
+ emo_jl = m.group("emo")
+ if emo_jl not in JL_MAP:
+ continue
+ label = JL_MAP[emo_jl]
+ txt_path = wav.with_suffix(".txt")
+ if not txt_path.exists():
+ continue
+ text = txt_path.read_text(encoding="utf-8", errors="ignore").strip()
+ if not text:
+ continue
+ rows.append({
+ "path": str(wav),
+ "text": text,
+ "label": label,
+ "source": "jl_corpus",
+ "speaker": m.group("spk"),
+ })
+ logger.info("JL-Corpus 5-class: %d samples", len(rows))
+ return rows
+
+
+def load_savee() -> list[dict]:
+ savee_dir = Path("data/savee/ALL")
+ asr_path = Path("data/savee/savee_asr.json")
+ if not asr_path.exists():
+ logger.warning("SAVEE ASR not found at %s", asr_path)
+ return []
+ asr = json.loads(asr_path.read_text())
+
+ rows: list[dict] = []
+ for wav in sorted(savee_dir.iterdir()):
+ m = SAVEE_PATTERN.match(wav.name)
+ if not m:
+ continue
+ label = SAVEE_MAP[m.group("emo")]
+ text = asr.get(wav.name, "").strip()
+ if not text:
+ continue
+ rows.append({
+ "path": str(wav),
+ "text": text,
+ "label": label,
+ "source": "savee",
+ "speaker": m.group("spk"),
+ })
+ logger.info("SAVEE disgust+surprise: %d samples", len(rows))
+ return rows
+
+
+def load_meld_disgust_surprise() -> list[dict]:
+ meld_path = Path("data/meld_fusion/manifest.json")
+ if not meld_path.exists():
+ logger.warning("MELD manifest missing: %s", meld_path)
+ return []
+ data = json.loads(meld_path.read_text())
+ rows: list[dict] = []
+ for d in data:
+ if d["label"] not in ("disgust", "surprise"):
+ continue
+ if not d.get("text", "").strip():
+ continue
+ rows.append({
+ "path": d["path"],
+ "text": d["text"],
+ "label": d["label"],
+ "source": "meld",
+ "speaker": f"dia{d.get('dialogue_id', '?')}",
+ })
+ logger.info("MELD disgust+surprise: %d samples", len(rows))
+ return rows
+
+
+def main() -> None:
+ OUT_DIR.mkdir(parents=True, exist_ok=True)
+ all_rows = load_jl_corpus() + load_savee() + load_meld_disgust_surprise()
+
+ counts = Counter(r["label"] for r in all_rows)
+ sources = Counter(r["source"] for r in all_rows)
+ logger.info("Total: %d samples", len(all_rows))
+ logger.info("By label: %s", dict(counts))
+ logger.info("By source: %s", dict(sources))
+
+ OUT_MANIFEST.write_text(json.dumps(all_rows, indent=2, ensure_ascii=False))
+ logger.info("Saved to %s", OUT_MANIFEST)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/build_meld_test_sets.py b/scripts/build_meld_test_sets.py
new file mode 100644
index 0000000000000000000000000000000000000000..a4df041aed92d4e0e58d1f3e9bc13cc1395e41ca
--- /dev/null
+++ b/scripts/build_meld_test_sets.py
@@ -0,0 +1,373 @@
+#!/usr/bin/env python3
+"""
+Build English test sets from MELD (Friends) dataset.
+
+Extracts 8 scenario-based test sets from MELD MP4 clips,
+converts to WAV, and concatenates into single audio files
+that simulate real phone calls for E2E pipeline testing.
+
+Usage:
+ python scripts/build_meld_test_sets.py
+
+Output:
+ data/meld_test/
+ โโโ 01_angry_fight.wav
+ โโโ 02_happy_loving.wav
+ โโโ ...
+ โโโ 08_calm_daily.wav
+ โโโ ground_truth.json # per-utterance emotion labels
+ โโโ README.md # test set descriptions
+"""
+
+import csv
+import json
+import os
+import subprocess
+import sys
+import tempfile
+from collections import Counter
+from pathlib import Path
+
+# --- Configuration ---
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+ZIP_PATH = PROJECT_ROOT / "data" / "english_test.zip"
+OUTPUT_DIR = PROJECT_ROOT / "data" / "meld_test"
+SAMPLE_RATE = 16000 # 16kHz mono โ matches our pipeline input
+
+# 8 test scenarios โ each maps to a specific MELD dialogue
+TEST_SETS = [
+ {
+ "tag": "01_angry_fight",
+ "desc": "Ross-Rachel breakup fight โ anger dominant (S3E15)",
+ "scenario": "Couple in a heated argument",
+ "primary_emotion": "anger",
+ "split": "train",
+ "dia_id": "51",
+ },
+ {
+ "tag": "02_happy_loving",
+ "desc": "Monica-Chandler sweet moment โ joy dominant (S5E14)",
+ "scenario": "Couple being affectionate and playful",
+ "primary_emotion": "joy",
+ "split": "train",
+ "dia_id": "1026",
+ },
+ {
+ "tag": "03_sad_emotional",
+ "desc": "Ross-Rachel emotional confession โ sadness dominant (S3E25)",
+ "scenario": "Emotional conversation with sadness and regret",
+ "primary_emotion": "sadness",
+ "split": "train",
+ "dia_id": "312",
+ },
+ {
+ "tag": "04_surprise_shock",
+ "desc": "Ross-Rachel surprise revelations (S7E18)",
+ "scenario": "Unexpected news and reactions",
+ "primary_emotion": "surprise",
+ "split": "train",
+ "dia_id": "747",
+ },
+ {
+ "tag": "05_fear_anxiety",
+ "desc": "Monica-Chandler anxious situation โ fear+mixed (S4E14)",
+ "scenario": "Anxious and worried conversation",
+ "primary_emotion": "fear",
+ "split": "train",
+ "dia_id": "109",
+ },
+ {
+ "tag": "06_disgust_annoyance",
+ "desc": "Family annoyance scene โ disgust+anger (S6E9)",
+ "scenario": "Annoyed and disgusted reactions",
+ "primary_emotion": "disgust",
+ "split": "train",
+ "dia_id": "1025",
+ },
+ {
+ "tag": "07_bittersweet",
+ "desc": "Ross-Rachel bittersweet farewell โ sadness+surprise (S5E5)",
+ "scenario": "Mixed emotions: saying goodbye with conflicting feelings",
+ "primary_emotion": "sadness",
+ "split": "train",
+ "dia_id": "676",
+ },
+ {
+ "tag": "08_calm_daily",
+ "desc": "Casual daily conversation โ neutral baseline (S3E23)",
+ "scenario": "Normal everyday chitchat (baseline)",
+ "primary_emotion": "neutral",
+ "split": "train",
+ "dia_id": "450",
+ },
+]
+
+
+def load_csv_from_zip(zip_path: Path) -> dict[str, list[dict]]:
+ """Load all CSV data from zip, grouped by split_diaID."""
+ import zipfile
+
+ dialogues = {}
+ with zipfile.ZipFile(zip_path, "r") as zf:
+ csv_files = [
+ ("train", "JSON files/JSON files/CSV Processed/train_sent_emo_cleaned_processed.csv"),
+ ("dev", "JSON files/JSON files/CSV Processed/dev_sent_emo_cleaned_processed.csv"),
+ ("test", "JSON files/JSON files/CSV Processed/test_sent_emo_cleaned_processed.csv"),
+ ]
+ for split, csv_path in csv_files:
+ try:
+ with zf.open(csv_path) as f:
+ import io
+ reader = csv.DictReader(io.TextIOWrapper(f, encoding="utf-8"))
+ for row in reader:
+ key = f"{split}_{row['Dialogue_ID']}"
+ dialogues.setdefault(key, []).append(row)
+ except KeyError:
+ print(f" Warning: {csv_path} not found in zip")
+ return dialogues
+
+
+def find_mp4_path(split: str, dia_id: str, utt_id: str, available_files: set) -> str | None:
+ """Find MP4 file path for a specific utterance."""
+ patterns = [
+ f"MELD.Raw/MELD.Raw/{split}/{split}_splits/dia{dia_id}_utt{utt_id}.mp4",
+ f"MELD.Raw/MELD.Raw/{split}/{split}_splits_complete/dia{dia_id}_utt{utt_id}.mp4",
+ f"MELD.Raw/MELD.Raw/{split}/output_repeated_splits_{split}/final_videos_{split}dia{dia_id}_utt{utt_id}.mp4",
+ ]
+ for p in patterns:
+ if p in available_files:
+ return p
+ return None
+
+
+def get_mp4_list_from_zip(zip_path: Path) -> set:
+ """Get set of all MP4 file paths in zip."""
+ import zipfile
+ with zipfile.ZipFile(zip_path, "r") as zf:
+ return {n for n in zf.namelist() if n.endswith(".mp4")}
+
+
+def extract_and_concat_wav(
+ zip_path: Path, mp4_paths: list[str], output_wav: Path, sample_rate: int = 16000
+) -> float:
+ """Extract audio from MP4s in zip and concatenate into single WAV."""
+ with tempfile.TemporaryDirectory() as tmpdir:
+ tmpdir = Path(tmpdir)
+ wav_parts = []
+
+ # Extract each MP4 and convert to WAV
+ import zipfile
+ with zipfile.ZipFile(zip_path, "r") as zf:
+ for i, mp4_path in enumerate(mp4_paths):
+ mp4_local = tmpdir / f"part_{i:03d}.mp4"
+ wav_local = tmpdir / f"part_{i:03d}.wav"
+
+ # Extract MP4
+ with zf.open(mp4_path) as src, open(mp4_local, "wb") as dst:
+ dst.write(src.read())
+
+ # Convert to WAV (16kHz mono)
+ result = subprocess.run(
+ [
+ "ffmpeg", "-y", "-i", str(mp4_local),
+ "-ar", str(sample_rate),
+ "-ac", "1",
+ "-acodec", "pcm_s16le",
+ str(wav_local),
+ ],
+ capture_output=True,
+ text=True,
+ )
+ if result.returncode != 0:
+ print(f" Warning: ffmpeg failed for {mp4_path}: {result.stderr[:200]}")
+ continue
+
+ if wav_local.exists() and wav_local.stat().st_size > 0:
+ wav_parts.append(wav_local)
+
+ if not wav_parts:
+ return 0.0
+
+ # Concatenate WAVs using ffmpeg concat
+ list_file = tmpdir / "concat_list.txt"
+ with open(list_file, "w") as f:
+ for wp in wav_parts:
+ f.write(f"file '{wp}'\n")
+
+ output_wav.parent.mkdir(parents=True, exist_ok=True)
+ result = subprocess.run(
+ [
+ "ffmpeg", "-y", "-f", "concat", "-safe", "0",
+ "-i", str(list_file),
+ "-ar", str(sample_rate),
+ "-ac", "1",
+ "-acodec", "pcm_s16le",
+ str(output_wav),
+ ],
+ capture_output=True,
+ text=True,
+ )
+ if result.returncode != 0:
+ print(f" Concat failed: {result.stderr[:300]}")
+ return 0.0
+
+ # Get duration
+ probe = subprocess.run(
+ ["ffprobe", "-v", "quiet", "-show_entries", "format=duration",
+ "-of", "default=noprint_wrappers=1:nokey=1", str(output_wav)],
+ capture_output=True, text=True,
+ )
+ try:
+ return float(probe.stdout.strip())
+ except ValueError:
+ return 0.0
+
+
+def main():
+ print("=" * 60)
+ print(" MELD English Test Set Builder")
+ print("=" * 60)
+
+ if not ZIP_PATH.exists():
+ print(f"Error: {ZIP_PATH} not found")
+ sys.exit(1)
+
+ # 1. Load CSV data
+ print("\n[1/4] Loading CSV data from zip...")
+ dialogues = load_csv_from_zip(ZIP_PATH)
+ print(f" Loaded {len(dialogues)} dialogues")
+
+ # 2. Get available MP4 files
+ print("[2/4] Scanning MP4 files in zip...")
+ mp4_files = get_mp4_list_from_zip(ZIP_PATH)
+ print(f" Found {len(mp4_files)} MP4 files")
+
+ # 3. Process each test set
+ print("[3/4] Building test sets...\n")
+ ground_truth = {}
+ summary_lines = []
+
+ for ts in TEST_SETS:
+ tag = ts["tag"]
+ key = f"{ts['split']}_{ts['dia_id']}"
+ utts = dialogues.get(key, [])
+
+ if not utts:
+ print(f" โ {tag}: dialogue {key} not found")
+ continue
+
+ print(f" ๐ฆ {tag} โ {ts['desc']}")
+ print(f" {len(utts)} utterances", end="")
+
+ # Find MP4 paths
+ mp4_paths = []
+ for u in utts:
+ p = find_mp4_path(ts["split"], ts["dia_id"], u["Utterance_ID"], mp4_files)
+ if p:
+ mp4_paths.append(p)
+
+ print(f", {len(mp4_paths)}/{len(utts)} MP4s found")
+
+ if not mp4_paths:
+ print(f" โ No MP4 files found, skipping")
+ continue
+
+ # Extract and concatenate
+ output_wav = OUTPUT_DIR / f"{tag}.wav"
+ duration = extract_and_concat_wav(ZIP_PATH, mp4_paths, output_wav, SAMPLE_RATE)
+ print(f" โ
{output_wav.name} โ {duration:.1f}s")
+
+ # Build ground truth
+ emo_counts = Counter(u["Emotion"] for u in utts)
+ ground_truth[tag] = {
+ "description": ts["desc"],
+ "scenario": ts["scenario"],
+ "primary_emotion": ts["primary_emotion"],
+ "source": f"MELD Friends S{utts[0]['Season']}E{utts[0]['Episode']} Dialogue {ts['dia_id']}",
+ "duration_sec": round(duration, 1),
+ "emotion_distribution": dict(emo_counts),
+ "total_utterances": len(utts),
+ "utterances": [
+ {
+ "speaker": u["Speaker"],
+ "emotion": u["Emotion"],
+ "sentiment": u["Sentiment"],
+ "text": u["Utterance"],
+ }
+ for u in utts
+ ],
+ }
+
+ summary_lines.append(
+ f"| {tag} | {ts['scenario'][:40]} | {ts['primary_emotion']} | {duration:.1f}s | {len(utts)} utts | {dict(emo_counts)} |"
+ )
+
+ # 4. Save ground truth + README
+ print("\n[4/4] Saving metadata...")
+
+ gt_path = OUTPUT_DIR / "ground_truth.json"
+ with open(gt_path, "w", encoding="utf-8") as f:
+ json.dump(ground_truth, f, indent=2, ensure_ascii=False)
+ print(f" โ
{gt_path}")
+
+ # Emotion alignment check
+ our_labels = {"neutral", "joy", "sadness", "anger", "surprise", "fear", "disgust"}
+ meld_labels = set()
+ for gt in ground_truth.values():
+ meld_labels.update(gt["emotion_distribution"].keys())
+
+ readme_content = f"""# MELD English Test Sets
+
+## Emotion Label Alignment
+
+| UsTwo Pipeline (EN) | MELD Label | Match |
+|---|---|---|
+| neutral | neutral | โ
Exact |
+| joy | joy | โ
Exact |
+| sadness | sadness | โ
Exact |
+| anger | anger | โ
Exact |
+| surprise | surprise | โ
Exact |
+| fear | fear | โ
Exact |
+| disgust | disgust | โ
Exact |
+
+**7/7 labels match exactly.** No mapping needed.
+
+## Test Sets
+
+| File | Scenario | Primary Emotion | Duration | Utterances | Emotion Distribution |
+|---|---|---|---|---|---|
+{chr(10).join(summary_lines)}
+
+## Source
+- Dataset: MELD (Multimodal EmotionLines Dataset)
+- Source: Friends TV series
+- Paper: Poria et al., ACL 2019
+- Each WAV is a full dialogue concatenated from per-utterance MP4 clips
+- Audio: 16kHz mono PCM (matches pipeline input format)
+
+## Usage
+```bash
+# Run pipeline on a single test set
+python scripts/run_pipeline.py data/meld_test/01_angry_fight.wav
+
+# Evaluate all test sets
+python scripts/evaluate_meld_test.py
+```
+"""
+ readme_path = OUTPUT_DIR / "README.md"
+ with open(readme_path, "w", encoding="utf-8") as f:
+ f.write(readme_content)
+ print(f" โ
{readme_path}")
+
+ # Summary
+ print("\n" + "=" * 60)
+ print(" DONE")
+ print("=" * 60)
+ total_files = len(list(OUTPUT_DIR.glob("*.wav")))
+ print(f" {total_files} WAV files in {OUTPUT_DIR}")
+ print(f" Ground truth: {gt_path}")
+ print(f" Emotion alignment: {len(our_labels & meld_labels)}/{len(our_labels)} exact match")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/cache_models.py b/scripts/cache_models.py
new file mode 100644
index 0000000000000000000000000000000000000000..be6ad227f44787e080102e96417608d59b8c2236
--- /dev/null
+++ b/scripts/cache_models.py
@@ -0,0 +1,128 @@
+#!/usr/bin/env python3
+"""Pre-download and cache all ML models at Docker build time.
+
+This avoids cold-start model downloads on first request in HF Spaces.
+Models are cached to default HuggingFace/torch hub directories.
+
+Usage (in Dockerfile):
+ RUN python scripts/cache_models.py
+"""
+
+import logging
+import os
+import sys
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
+logger = logging.getLogger("cache_models")
+
+
+def cache_pyannote():
+ """Cache pyannote speaker-diarization-3.1 (~1.5GB)."""
+ hf_token = os.environ.get("HF_TOKEN")
+ if not hf_token:
+ logger.warning("HF_TOKEN not set, skipping pyannote cache")
+ return
+ try:
+ from pyannote.audio import Pipeline
+ logger.info("Caching pyannote/speaker-diarization-3.1...")
+ Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", token=hf_token)
+ logger.info("pyannote cached OK")
+ except Exception as e:
+ logger.warning("pyannote cache failed: %s", e)
+
+
+def cache_whisperx():
+ """Cache WhisperX large-v3-turbo INT8 (~1.5GB)."""
+ try:
+ import whisperx
+ logger.info("Caching whisperx large-v3-turbo (int8)...")
+ whisperx.load_model("large-v3-turbo", device="cpu", compute_type="int8")
+ logger.info("whisperx cached OK")
+ except Exception as e:
+ logger.warning("whisperx cache failed: %s", e)
+
+
+def cache_emotion2vec():
+ """Cache emotion2vec_plus_base (~300MB)."""
+ try:
+ from funasr import AutoModel
+ logger.info("Caching iic/emotion2vec_plus_base...")
+ AutoModel(model="iic/emotion2vec_plus_base", device="cpu", hub="hf")
+ logger.info("emotion2vec cached OK")
+ except Exception as e:
+ logger.warning("emotion2vec cache failed: %s", e)
+
+
+def cache_text_models():
+ """Cache text emotion models (~300MB each)."""
+ try:
+ from transformers import pipeline
+ logger.info("Caching j-hartmann/emotion-english-distilroberta-base...")
+ pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", top_k=None)
+ logger.info("DistilRoBERTa cached OK")
+ except Exception as e:
+ logger.warning("DistilRoBERTa cache failed: %s", e)
+
+ try:
+ from transformers import pipeline
+ logger.info("Caching searle-j/kote_for_easygoing_people...")
+ pipeline("text-classification", model="searle-j/kote_for_easygoing_people", top_k=None)
+ logger.info("KcELECTRA cached OK")
+ except Exception as e:
+ logger.warning("KcELECTRA cache failed: %s", e)
+
+
+def download_lora_onnx_models():
+ """Download LoRA ONNX models from public HF Hub repo into data/models/."""
+ from pathlib import Path
+ from huggingface_hub import hf_hub_download, snapshot_download
+
+ repo_id = "BBBAKERY/ustwo-lora-models"
+
+ # emotion2vec ONNX โ data/models/lora_emotion2vec_7class/model.onnx
+ audio_dir = Path("data/models/lora_emotion2vec_7class")
+ audio_dir.mkdir(parents=True, exist_ok=True)
+ try:
+ logger.info("Downloading emotion2vec LoRA ONNX from %s...", repo_id)
+ for fname in ["emotion2vec/model.onnx", "emotion2vec/model.json"]:
+ path = hf_hub_download(repo_id=repo_id, filename=fname, repo_type="model")
+ target = audio_dir / Path(fname).name
+ if not target.exists() or target.resolve() != Path(path).resolve():
+ import shutil
+ shutil.copy(path, target)
+ logger.info("emotion2vec LoRA ONNX cached OK")
+ except Exception as e:
+ logger.warning("emotion2vec LoRA ONNX download failed: %s", e)
+
+ # KcELECTRA ONNX + tokenizer โ data/models/lora_kcelectra_7class/
+ text_dir = Path("data/models/lora_kcelectra_7class")
+ text_dir.mkdir(parents=True, exist_ok=True)
+ try:
+ logger.info("Downloading KcELECTRA LoRA ONNX from %s...", repo_id)
+ for fname in ["kcelectra/model.onnx", "kcelectra/model.json"]:
+ path = hf_hub_download(repo_id=repo_id, filename=fname, repo_type="model")
+ target = text_dir / Path(fname).name
+ import shutil
+ shutil.copy(path, target)
+
+ # Tokenizer folder
+ tokenizer_target = text_dir / "best_model"
+ tokenizer_target.mkdir(parents=True, exist_ok=True)
+ for fname in ["tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "vocab.txt"]:
+ path = hf_hub_download(repo_id=repo_id, filename=f"kcelectra/tokenizer/{fname}", repo_type="model")
+ import shutil
+ shutil.copy(path, tokenizer_target / fname)
+
+ logger.info("KcELECTRA LoRA ONNX + tokenizer cached OK")
+ except Exception as e:
+ logger.warning("KcELECTRA LoRA ONNX download failed: %s", e)
+
+
+if __name__ == "__main__":
+ logger.info("=== Pre-caching ML models ===")
+ cache_pyannote()
+ cache_whisperx()
+ cache_emotion2vec()
+ cache_text_models()
+ download_lora_onnx_models()
+ logger.info("=== Model caching complete ===")
diff --git a/scripts/convert_to_onnx.py b/scripts/convert_to_onnx.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/scripts/eval_audio_on_subset.py b/scripts/eval_audio_on_subset.py
new file mode 100644
index 0000000000000000000000000000000000000000..d47761f29d03263a87313fd58a385223e50e1299
--- /dev/null
+++ b/scripts/eval_audio_on_subset.py
@@ -0,0 +1,138 @@
+#!/usr/bin/env python3
+"""Evaluate audio model (LoRA ONNX or base) on a subset of val manifest.
+
+Usage:
+ python scripts/eval_audio_on_subset.py \
+ --val-manifest data/lora_dataset/val_manifest.json \
+ --source ravdess \
+ --model lora_onnx \
+ --onnx data/models/lora_emotion2vec_7class/model.onnx
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+from collections import Counter
+from pathlib import Path
+
+import numpy as np
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+PROJECT_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear", "disgust"]
+
+LORA_LABELS = ["happiness", "anger", "disgust", "fear", "neutral", "sadness", "surprise"]
+LORA_TO_PROJECT = {
+ "happiness": "joy", "anger": "anger", "disgust": "disgust",
+ "fear": "fear", "neutral": "neutral", "sadness": "sadness", "surprise": "surprise",
+}
+
+BASE_LABEL_MAP = {
+ "angry": "anger", "disgusted": "disgust", "fearful": "fear",
+ "happy": "joy", "neutral": "neutral", "sad": "sadness", "surprised": "surprise",
+ "other": "neutral", "unknown": "neutral",
+ "็ๆฐ/angry": "anger", "ๅๆถ/disgusted": "disgust", "ๆๆง/fearful": "fear",
+ "ๅผๅฟ/happy": "joy", "ไธญ็ซ/neutral": "neutral", "้พ่ฟ/sad": "sadness",
+ "ๅๆ/surprised": "surprise", "ๅ
ถไป/other": "neutral", "": "neutral",
+}
+
+
+def predict_lora_onnx(audio_path: str, session, max_seconds: float = 15.0):
+ import soundfile as sf
+ audio, sr = sf.read(audio_path, dtype="float32")
+ if audio.ndim == 2:
+ audio = audio.mean(axis=1)
+ if sr != 16000:
+ import librosa
+ audio = librosa.resample(audio, orig_sr=sr, target_sr=16000)
+ max_samples = int(max_seconds * 16000)
+ if len(audio) > max_samples:
+ audio = audio[:max_samples]
+ waveform = audio.reshape(1, -1).astype(np.float32)
+ logits = session.run(None, {"waveform": waveform})[0]
+ exp_logits = np.exp(logits - logits.max(axis=-1, keepdims=True))
+ probs = (exp_logits / exp_logits.sum(axis=-1, keepdims=True)).squeeze()
+ scores = {label: 0.0 for label in PROJECT_LABELS}
+ for lora_label, prob in zip(LORA_LABELS, probs):
+ scores[LORA_TO_PROJECT[lora_label]] = float(prob)
+ return max(scores, key=scores.get)
+
+
+def predict_base(audio_path: str, funasr_model):
+ try:
+ output = funasr_model.generate(audio_path, granularity="utterance", extract_embedding=False)
+ except Exception:
+ return "neutral"
+ scores = {label: 0.0 for label in PROJECT_LABELS}
+ if output and isinstance(output, list) and len(output) > 0:
+ rec = output[0]
+ for native_label, score in zip(rec.get("labels", []), rec.get("scores", [])):
+ pl = BASE_LABEL_MAP.get(native_label, "neutral")
+ scores[pl] += float(score)
+ return max(scores, key=scores.get)
+
+
+def f1_score(y_true, y_pred, label):
+ tp = sum(1 for t, p in zip(y_true, y_pred) if t == label and p == label)
+ fp = sum(1 for t, p in zip(y_true, y_pred) if t != label and p == label)
+ fn = sum(1 for t, p in zip(y_true, y_pred) if t == label and p != label)
+ if tp + fp == 0 or tp + fn == 0:
+ return 0.0
+ p = tp / (tp + fp); r = tp / (tp + fn)
+ return 2 * p * r / (p + r) if (p + r) > 0 else 0.0
+
+
+def main():
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--val-manifest", default="data/lora_dataset/val_manifest.json")
+ parser.add_argument("--source", default="ravdess", help="Filter by source: ravdess, 263, 71631")
+ parser.add_argument("--model", choices=["lora_onnx", "base"], default="lora_onnx")
+ parser.add_argument("--onnx", default="data/models/lora_emotion2vec_7class/model.onnx")
+ args = parser.parse_args()
+
+ with open(args.val_manifest) as f:
+ val = json.load(f)
+ samples = [s for s in val if s["source"] == args.source]
+ logger.info("Evaluating %s on %d %s samples", args.model, len(samples), args.source)
+
+ # Normalize labels: happiness โ joy
+ for s in samples:
+ if s["label"] == "happiness":
+ s["label"] = "joy"
+
+ if args.model == "lora_onnx":
+ import onnxruntime as ort
+ session = ort.InferenceSession(args.onnx, providers=["CPUExecutionProvider"])
+ predict_fn = lambda p: predict_lora_onnx(p, session)
+ else:
+ from funasr import AutoModel
+ model = AutoModel(model="iic/emotion2vec_plus_base", device="cpu", hub="hf")
+ predict_fn = lambda p: predict_base(p, model)
+
+ y_true, y_pred = [], []
+ for i, s in enumerate(samples):
+ y_true.append(s["label"])
+ y_pred.append(predict_fn(s["path"]))
+ if (i + 1) % 100 == 0:
+ logger.info("Progress: %d / %d", i + 1, len(samples))
+
+ # Per-class F1
+ f1s = {label: f1_score(y_true, y_pred, label) for label in PROJECT_LABELS}
+ macro_f1 = np.mean(list(f1s.values()))
+ acc = sum(1 for t, p in zip(y_true, y_pred) if t == p) / len(samples)
+
+ print()
+ print(f"=== {args.model} on {args.source} ({len(samples)} samples) ===")
+ print(f"Macro F1: {macro_f1:.4f}")
+ print(f"Accuracy: {acc:.4f}")
+ print("Per-class F1:")
+ for label, f1 in f1s.items():
+ support = sum(1 for t in y_true if t == label)
+ if support > 0:
+ print(f" {label:<12} {f1:.4f} (n={support})")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/evaluate_emotion2vec_english.py b/scripts/evaluate_emotion2vec_english.py
new file mode 100644
index 0000000000000000000000000000000000000000..df20c66644ade958bdc84772071080c12770ca41
--- /dev/null
+++ b/scripts/evaluate_emotion2vec_english.py
@@ -0,0 +1,194 @@
+#!/usr/bin/env python3
+"""emotion2vec base ๋ชจ๋ธ ์์ด ํ๊ฐ (RAVDESS ๋ฐ์ดํฐ์
).
+
+clean + phone 2๊ฐ ์กฐ๊ฑด์ผ๋ก ํ๊ฐํ์ฌ ์ค์ ํตํ ํ๊ฒฝ ์ฑ๋ฅ์ ์ถ์ ํ๋ค.
+
+Usage:
+ python scripts/evaluate_emotion2vec_english.py
+ python scripts/evaluate_emotion2vec_english.py --condition clean # clean๋ง
+ python scripts/evaluate_emotion2vec_english.py --condition phone # phone๋ง
+ python scripts/evaluate_emotion2vec_english.py --max-samples 100 # ๋น ๋ฅธ ํ
์คํธ
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import logging
+import sys
+import time
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).parent.parent
+sys.path.insert(0, str(PROJECT_ROOT))
+
+logging.basicConfig(
+ level=logging.INFO,
+ format="%(asctime)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger("eval_emotion2vec_en")
+
+MANIFEST_PATH = PROJECT_ROOT / "data" / "ravdess" / "manifest.csv"
+OUTPUT_JSON = PROJECT_ROOT / "data" / "ravdess_eval_results.json"
+
+PROJECT_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear", "disgust"]
+
+
+def load_manifest(max_samples: int | None = None) -> list[dict]:
+ """manifest.csv ๋ก๋."""
+ rows = []
+ with open(MANIFEST_PATH) as f:
+ reader = csv.DictReader(f)
+ for row in reader:
+ rows.append(row)
+ if max_samples:
+ rows = rows[:max_samples]
+ logger.info(f"manifest ๋ก๋: {len(rows)}๊ฐ ์ํ")
+ return rows
+
+
+def evaluate_condition(
+ samples: list[dict],
+ condition: str,
+ device: str,
+) -> dict:
+ """ํ ์กฐ๊ฑด(clean/phone)์ ๋ํด ํ๊ฐ ์คํ."""
+ from sklearn.metrics import (
+ accuracy_score,
+ classification_report,
+ confusion_matrix,
+ )
+ from src.stage2.audio_emotion import predict as audio_predict
+
+ path_key = "clean_path" if condition == "clean" else "phone_path"
+
+ y_true = []
+ y_pred = []
+ latencies = []
+ errors = 0
+
+ total = len(samples)
+ for i, sample in enumerate(samples, 1):
+ audio_path = sample[path_key]
+ if not audio_path or not Path(audio_path).exists():
+ errors += 1
+ continue
+
+ ground_truth = sample["emotion"]
+
+ t0 = time.perf_counter()
+ result = audio_predict(audio_path, device=device)
+ latency = (time.perf_counter() - t0) * 1000 # ms
+
+ y_true.append(ground_truth)
+ y_pred.append(result["emotion"])
+ latencies.append(latency)
+
+ if i % 200 == 0 or i == total:
+ acc_so_far = sum(1 for t, p in zip(y_true, y_pred) if t == p) / len(y_true)
+ logger.info(
+ f" [{condition}] {i}/{total} โ "
+ f"acc={acc_so_far:.3f}, "
+ f"avg_latency={sum(latencies)/len(latencies):.0f}ms"
+ )
+
+ accuracy = accuracy_score(y_true, y_pred)
+ report = classification_report(
+ y_true, y_pred, labels=PROJECT_LABELS, output_dict=True, zero_division=0,
+ )
+ cm = confusion_matrix(y_true, y_pred, labels=PROJECT_LABELS)
+
+ # per-class metrics ์ ๋ฆฌ
+ per_class = {}
+ for label in PROJECT_LABELS:
+ if label in report:
+ per_class[label] = {
+ "precision": round(report[label]["precision"], 4),
+ "recall": round(report[label]["recall"], 4),
+ "f1": round(report[label]["f1-score"], 4),
+ "support": report[label]["support"],
+ }
+
+ result = {
+ "condition": condition,
+ "total_samples": len(y_true),
+ "errors": errors,
+ "accuracy": round(accuracy, 4),
+ "macro_f1": round(report["macro avg"]["f1-score"], 4),
+ "weighted_f1": round(report["weighted avg"]["f1-score"], 4),
+ "per_class": per_class,
+ "confusion_matrix": cm.tolist(),
+ "confusion_labels": PROJECT_LABELS,
+ "avg_latency_ms": round(sum(latencies) / len(latencies), 1) if latencies else 0,
+ }
+
+ logger.info(f"\n{'='*60}")
+ logger.info(f"[{condition.upper()}] ๊ฒฐ๊ณผ:")
+ logger.info(f" Accuracy: {accuracy:.4f}")
+ logger.info(f" Macro F1: {report['macro avg']['f1-score']:.4f}")
+ logger.info(f" Weighted F1: {report['weighted avg']['f1-score']:.4f}")
+ logger.info(f" Avg Latency: {result['avg_latency_ms']:.0f}ms")
+ logger.info(f"\nPer-class F1:")
+ for label in PROJECT_LABELS:
+ if label in per_class:
+ logger.info(f" {label:10s}: F1={per_class[label]['f1']:.3f} "
+ f"(P={per_class[label]['precision']:.3f} R={per_class[label]['recall']:.3f}) "
+ f"n={per_class[label]['support']}")
+ logger.info(f"\nConfusion Matrix (rows=true, cols=pred):")
+ logger.info(f" {'':10s} " + " ".join(f"{l[:4]:>6s}" for l in PROJECT_LABELS))
+ for i_row, label in enumerate(PROJECT_LABELS):
+ row_str = " ".join(f"{v:6d}" for v in cm[i_row])
+ logger.info(f" {label:10s} {row_str}")
+ logger.info(f"{'='*60}\n")
+
+ return result
+
+
+def main():
+ parser = argparse.ArgumentParser(description="emotion2vec ์์ด ํ๊ฐ (RAVDESS)")
+ parser.add_argument("--condition", choices=["clean", "phone", "both"], default="both")
+ parser.add_argument("--device", default="cpu")
+ parser.add_argument("--max-samples", type=int, default=None, help="ํ๊ฐํ ์ต๋ ์ํ ์")
+ args = parser.parse_args()
+
+ if not MANIFEST_PATH.exists():
+ logger.error(f"manifest.csv๋ฅผ ์ฐพ์ ์ ์์ต๋๋ค. ๋จผ์ prepare_ravdess.py๋ฅผ ์คํํ์ธ์.")
+ sys.exit(1)
+
+ samples = load_manifest(args.max_samples)
+
+ conditions = []
+ if args.condition in ("clean", "both"):
+ conditions.append("clean")
+ if args.condition in ("phone", "both"):
+ conditions.append("phone")
+
+ results = []
+ for cond in conditions:
+ logger.info(f"\n{'#'*60}")
+ logger.info(f"ํ๊ฐ ์์: {cond.upper()} ์กฐ๊ฑด")
+ logger.info(f"{'#'*60}")
+ result = evaluate_condition(samples, cond, args.device)
+ results.append(result)
+
+ # ๊ฒฐ๊ณผ ์ ์ฅ
+ with open(OUTPUT_JSON, "w") as f:
+ json.dump(results, f, indent=2, ensure_ascii=False)
+ logger.info(f"๊ฒฐ๊ณผ ์ ์ฅ: {OUTPUT_JSON}")
+
+ # clean vs phone ๋น๊ต (both์ผ ๋)
+ if len(results) == 2:
+ clean_acc = results[0]["accuracy"]
+ phone_acc = results[1]["accuracy"]
+ degradation = clean_acc - phone_acc
+ logger.info(f"\n{'='*60}")
+ logger.info(f"Clean vs Phone ๋น๊ต:")
+ logger.info(f" Clean accuracy: {clean_acc:.4f}")
+ logger.info(f" Phone accuracy: {phone_acc:.4f}")
+ logger.info(f" Degradation: {degradation:+.4f}")
+ logger.info(f"{'='*60}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/export_lora_onnx.py b/scripts/export_lora_onnx.py
new file mode 100644
index 0000000000000000000000000000000000000000..550e1e916f96d110e688cad97e33b59dc8479c3d
--- /dev/null
+++ b/scripts/export_lora_onnx.py
@@ -0,0 +1,269 @@
+#!/usr/bin/env python3
+"""ONNX export for LoRA-finetuned emotion2vec 7-class model.
+
+Merges LoRA weights into base model, wraps as a single waveform-to-logits
+module, and exports to ONNX with dynamic batch/time axes.
+
+Usage:
+ python scripts/export_lora_onnx.py \
+ --checkpoint data/models/lora_emotion2vec_7class/best_lora.pt \
+ --output data/models/lora_emotion2vec_7class/emotion2vec_lora.onnx \
+ --device cpu
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+# Import LoRA components
+import sys
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+
+from train_lora_emotion2vec import (
+ LoRALinear,
+ MLPHead,
+ inject_lora,
+ merge_lora_linear,
+ LABELS_7CLASS,
+ NUM_CLASSES,
+)
+
+
+def merge_all_lora(encoder: nn.Module) -> None:
+ """Walk encoder.blocks, replace each LoRALinear with merged nn.Linear.
+
+ Modifies the encoder in-place.
+ """
+ for block in encoder.blocks:
+ if isinstance(block.attn.qkv, LoRALinear):
+ block.attn.qkv = merge_lora_linear(block.attn.qkv)
+ if isinstance(block.attn.proj, LoRALinear):
+ block.attn.proj = merge_lora_linear(block.attn.proj)
+
+
+class Emotion2vecONNXWrapper(nn.Module):
+ """Wraps emotion2vec encoder for ONNX export.
+
+ forward(waveform: (B, T)) -> logits: (B, 7)
+ Includes layer_norm + extract_features + mean pool + proj.
+ """
+
+ def __init__(self, encoder):
+ super().__init__()
+ self.encoder = encoder
+
+ def forward(self, waveform: torch.Tensor) -> torch.Tensor:
+ """
+ Args:
+ waveform: (B, T) float32, 16kHz
+
+ Returns:
+ logits: (B, 7)
+ """
+ # Layer norm (per-sample, ONNX-compatible: normalize along time axis)
+ if self.encoder.cfg.normalize:
+ mean = waveform.mean(dim=-1, keepdim=True)
+ var = waveform.var(dim=-1, keepdim=True, unbiased=False)
+ waveform = (waveform - mean) / torch.sqrt(var + 1e-5)
+
+ # Extract features
+ feats = self.encoder.extract_features(waveform, padding_mask=None)
+ x = feats["x"] # (B, T', 768)
+
+ # Mean pool
+ pooled = x.mean(dim=1) # (B, 768)
+
+ # Classify
+ logits = self.encoder.proj(pooled) # (B, 7)
+ return logits
+
+
+def export_onnx(checkpoint_path: str, output_path: str, device: str = "cpu"):
+ """Load base emotion2vec, inject LoRA, load checkpoint, merge, export ONNX.
+
+ Args:
+ checkpoint_path: path to LoRA checkpoint (.pt)
+ output_path: path for output ONNX file
+ device: "cpu" or "cuda"
+ """
+ from funasr import AutoModel
+
+ checkpoint_path = Path(checkpoint_path)
+ output_path = Path(output_path)
+ output_path.parent.mkdir(parents=True, exist_ok=True)
+
+ # Load checkpoint to get config
+ logger.info("Loading checkpoint: %s", checkpoint_path)
+ ckpt = torch.load(str(checkpoint_path), map_location=device, weights_only=True)
+
+ # Load base model
+ logger.info("Loading emotion2vec_plus_base...")
+ fmodel = AutoModel(model="iic/emotion2vec_plus_base", device=device, hub="hf")
+ encoder = fmodel.model
+
+ # Freeze + inject LoRA (dropout=0 for inference)
+ for param in encoder.parameters():
+ param.requires_grad = False
+ inject_lora(encoder, r=16, alpha=32, dropout=0.0)
+
+ # Replace proj with MLPHead
+ num_classes = ckpt.get("num_classes", NUM_CLASSES)
+ encoder.proj = MLPHead(768, num_classes, dropout=0.0).to(device)
+
+ # Load LoRA weights
+ lora_weights = ckpt["lora_weights"]
+ for name, module in encoder.named_modules():
+ if isinstance(module, LoRALinear):
+ a_key = f"{name}.lora_A.weight"
+ b_key = f"{name}.lora_B.weight"
+ if a_key in lora_weights:
+ module.lora_A.weight.data.copy_(lora_weights[a_key])
+ if b_key in lora_weights:
+ module.lora_B.weight.data.copy_(lora_weights[b_key])
+
+ # Load proj state
+ encoder.proj.load_state_dict(ckpt["proj"])
+
+ # Merge LoRA into base weights
+ logger.info("Merging LoRA weights...")
+ merge_all_lora(encoder)
+
+ # Verify no LoRALinear remains
+ lora_count = sum(1 for m in encoder.modules() if isinstance(m, LoRALinear))
+ assert lora_count == 0, f"Merge failed: {lora_count} LoRALinear remain"
+
+ # Wrap for ONNX
+ wrapper = Emotion2vecONNXWrapper(encoder)
+ wrapper.eval()
+
+ # Dummy input (1 second of audio at 16kHz)
+ dummy_input = torch.randn(1, 16000, device=device)
+
+ # Export
+ logger.info("Exporting ONNX to %s ...", output_path)
+ torch.onnx.export(
+ wrapper,
+ dummy_input,
+ str(output_path),
+ opset_version=17,
+ input_names=["waveform"],
+ output_names=["logits"],
+ dynamic_axes={
+ "waveform": {0: "batch", 1: "time"},
+ "logits": {0: "batch"},
+ },
+ )
+ logger.info("ONNX export complete: %s", output_path)
+
+ # Save label metadata JSON alongside
+ meta_path = output_path.with_suffix(".json")
+ meta = {
+ "model": "emotion2vec_plus_base + LoRA (merged)",
+ "num_classes": num_classes,
+ "labels": ckpt.get("labels", LABELS_7CLASS),
+ "input": "waveform: (batch, time) float32, 16kHz mono",
+ "output": f"logits: (batch, {num_classes}) float32",
+ "checkpoint": str(checkpoint_path),
+ }
+ with open(meta_path, "w") as f:
+ json.dump(meta, f, indent=2)
+ logger.info("Metadata saved: %s", meta_path)
+
+
+def export_kcelectra_onnx(model_dir: str, output_path: str, base_model_id: str = "beomi/KcELECTRA-base-v2022"):
+ """Export LoRA-finetuned KcELECTRA to ONNX.
+
+ Merges PEFT LoRA into base model, then exports text โ logits ONNX.
+ """
+ from transformers import AutoTokenizer, AutoModelForSequenceClassification
+ from peft import PeftModel
+
+ output_path = Path(output_path)
+ model_dir = Path(model_dir)
+
+ logger.info("Loading base model: %s", base_model_id)
+ base_model = AutoModelForSequenceClassification.from_pretrained(
+ base_model_id, num_labels=7,
+ )
+ tokenizer = AutoTokenizer.from_pretrained(str(model_dir))
+
+ logger.info("Loading PEFT adapter: %s", model_dir)
+ model = PeftModel.from_pretrained(base_model, str(model_dir))
+
+ logger.info("Merging LoRA weights...")
+ model = model.merge_and_unload()
+ model.eval()
+
+ # Dummy input
+ dummy = tokenizer("ํ
์คํธ ๋ฌธ์ฅ์
๋๋ค", return_tensors="pt", max_length=128, truncation=True, padding="max_length")
+
+ output_path.parent.mkdir(parents=True, exist_ok=True)
+ logger.info("Exporting ONNX to %s", output_path)
+
+ torch.onnx.export(
+ model,
+ (dummy["input_ids"], dummy["attention_mask"]),
+ str(output_path),
+ opset_version=17,
+ input_names=["input_ids", "attention_mask"],
+ output_names=["logits"],
+ dynamic_axes={
+ "input_ids": {0: "batch", 1: "seq_len"},
+ "attention_mask": {0: "batch", 1: "seq_len"},
+ "logits": {0: "batch"},
+ },
+ )
+
+ # Save metadata
+ labels = ["happiness", "anger", "disgust", "fear", "neutral", "sadness", "surprise"]
+ meta_path = output_path.with_suffix(".json")
+ meta = {
+ "model": f"{base_model_id} + LoRA (merged)",
+ "num_classes": 7,
+ "labels": labels,
+ "input": "input_ids: (batch, seq_len) int64, attention_mask: (batch, seq_len) int64",
+ "output": "logits: (batch, 7) float32",
+ "max_length": 128,
+ }
+ with open(meta_path, "w") as f:
+ json.dump(meta, f, indent=2, ensure_ascii=False)
+
+ logger.info("ONNX export complete: %s (+ %s)", output_path, meta_path)
+
+
+def main():
+ parser = argparse.ArgumentParser(description="Export LoRA models to ONNX")
+ parser.add_argument("--mode", required=True, choices=["audio", "text"],
+ help="audio: emotion2vec, text: KcELECTRA")
+ # Audio args
+ parser.add_argument("--checkpoint", help="Path to LoRA checkpoint (.pt) for audio mode")
+ # Text args
+ parser.add_argument("--model-dir", help="Path to PEFT model directory for text mode")
+ parser.add_argument("--base-model", default="beomi/KcELECTRA-base-v2022")
+ # Common
+ parser.add_argument("--output", required=True, help="Output ONNX path")
+ parser.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
+ args = parser.parse_args()
+
+ if args.mode == "audio":
+ if not args.checkpoint:
+ parser.error("--checkpoint required for audio mode")
+ export_onnx(args.checkpoint, args.output, args.device)
+ else:
+ if not args.model_dir:
+ parser.error("--model-dir required for text mode")
+ export_kcelectra_onnx(args.model_dir, args.output, args.base_model)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/optimize_fusion_weights.py b/scripts/optimize_fusion_weights.py
new file mode 100644
index 0000000000000000000000000000000000000000..73bf26fe7658f3d11900a7f07e78dbd0b1770585
--- /dev/null
+++ b/scripts/optimize_fusion_weights.py
@@ -0,0 +1,618 @@
+#!/usr/bin/env python3
+"""Grid Search for Emotion-Specific Fusion Weights.
+
+Uses AI Hub 263 val split (audio + text + ground truth) to find optimal
+audio/text fusion weights per emotion class.
+
+Outputs:
+ - fusion_grid_search.json โ full weight-F1 curves per emotion
+ - optimal_fusion_weights.json โ best weights per emotion
+ - fusion_grid_search.png โ 7 subplots: weight vs F1 per emotion
+ - fusion_comparison.png โ bar chart: fixed 60/40 vs optimal
+ - fusion_report.md โ text summary
+
+Usage:
+ python scripts/optimize_fusion_weights.py \
+ --val-manifest data/lora_dataset/val_manifest.json \
+ --onnx-model data/models/lora_emotion2vec_7class/model.onnx \
+ --anchor-dir "data/AI Hub 263" \
+ --output-dir data/models/lora_emotion2vec_7class
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import gc
+import json
+import logging
+import sys
+from collections import Counter, defaultdict
+from pathlib import Path
+
+import numpy as np
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+PROJECT_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear", "disgust"]
+
+# LoRA model labels โ project labels
+LORA_LABELS = ["happiness", "anger", "disgust", "fear", "neutral", "sadness", "surprise"]
+LORA_TO_PROJECT = {
+ "happiness": "joy", "anger": "anger", "disgust": "disgust",
+ "fear": "fear", "neutral": "neutral", "sadness": "sadness", "surprise": "surprise",
+}
+
+# 263 label mapping (same as prepare_lora_dataset.py)
+MAP_263 = {
+ "angry": "anger", "anger": "anger",
+ "sadness": "sadness", "sad": "sadness",
+ "happiness": "happiness", "happy": "happiness",
+ "fear": "fear", "disgust": "disgust",
+ "surprise": "surprise", "neutral": "neutral",
+}
+
+# KcELECTRA 44-class โ 7-class (from src/stage2/text_emotion.py)
+KO_LABEL_MAP = {
+ "๊ธฐ์จ": "joy", "์ฆ๊ฑฐ์/์ ๋จ": "joy", "ํ๋ณต": "joy",
+ "๊ฐ๋/๊ฐํ": "joy", "๊ณ ๋ง์": "joy", "ํ์/ํธ์": "joy",
+ "๋ฟ๋ฏํจ": "joy", "ํ๋ญํจ(๊ท์ฌ์/์์จ)": "joy", "๊ธฐ๋๊ฐ": "joy",
+ "ํธ์/์พ์ ": "joy", "์์ฌ/์ ๋ขฐ": "joy", "์๊ปด์ฃผ๋": "joy", "์กด๊ฒฝ": "joy",
+ "๋๋": "surprise", "์ ๊ธฐํจ/๊ด์ฌ": "surprise", "๊ฒฝ์
": "surprise", "์ด์ด์์": "surprise",
+ "์ฌํ": "sadness", "์๋ฌ์": "sadness", "์ํ๊น์/์ค๋ง": "sadness",
+ "์ ๋ง": "sadness", "๋ถ๋๋ฌ์": "sadness", "๋ถ์ํจ/์ฐ๋ฏผ": "sadness",
+ "ํจ๋ฐฐ/์๊ธฐํ์ค": "sadness", "ํ๋ฆ/์ง์นจ": "sadness", "์ฃ์ฑ
๊ฐ": "sadness",
+ "ํ๋จ/๋ถ๋
ธ": "anger", "์ง์ฆ": "anger", "๋ถํ/๋ถ๋ง": "anger",
+ "์ง๊ธ์ง๊ธ": "anger", "์ฐ์ญ๋/๋ฌด์ํจ": "anger", "ํ์ฌํจ": "anger",
+ "์ฆ์ค/ํ์ค": "anger", "๊ท์ฐฎ์": "anger",
+ "๊ณตํฌ/๋ฌด์์": "fear", "๋ถ์/๊ฑฑ์ ": "fear", "๋นํฉ/๋์ฒ": "fear", "์์ฌ/๋ถ์ ": "fear",
+ "์์": "neutral", "๊นจ๋ฌ์": "neutral", "์ฌ๋ฏธ์์": "neutral",
+ "๋ถ๋ด/์_๋ดํด": "neutral", "๋น์ฅํจ": "neutral",
+ "์ญ๊ฒจ์/์ง๊ทธ๋ฌ์": "disgust",
+}
+
+
+def load_263_texts(anchor_dir: Path) -> dict[str, str]:
+ """Load wav_id โ ๋ฐํ๋ฌธ mapping from 263 CSVs."""
+ texts = {}
+ for csv_path in sorted(anchor_dir.glob("*.csv")):
+ with open(csv_path, encoding="cp949") as f:
+ reader = csv.reader(f)
+ next(reader) # skip header
+ for row in reader:
+ wav_id = row[0]
+ text = row[1]
+ texts[wav_id] = text
+ logger.info("Loaded %d texts from 263 CSVs", len(texts))
+ return texts
+
+
+def predict_audio_base(audio_path: str, funasr_model, max_seconds: float = 15.0) -> dict[str, float]:
+ """Run base (non-finetuned) emotion2vec via FunASR, 9-class โ 7-class mapping.
+
+ Audio trimmed to max_seconds โ FunASR transformer has quadratic memory in sequence length,
+ so a 100s clip can blow past 15GB RAM. Matches predict_audio_onnx() behavior.
+ """
+ # emotion2vec base native labels โ project labels
+ LABEL_MAP = {
+ "angry": "anger", "disgusted": "disgust", "fearful": "fear",
+ "happy": "joy", "neutral": "neutral", "sad": "sadness", "surprised": "surprise",
+ "other": "neutral", "unknown": "neutral",
+ "็ๆฐ/angry": "anger", "ๅๆถ/disgusted": "disgust", "ๆๆง/fearful": "fear",
+ "ๅผๅฟ/happy": "joy", "ไธญ็ซ/neutral": "neutral", "้พ่ฟ/sad": "sadness",
+ "ๅๆ/surprised": "surprise", "ๅ
ถไป/other": "neutral", "": "neutral",
+ }
+
+ import soundfile as sf
+
+ audio, sr = sf.read(audio_path, dtype="float32")
+ if audio.ndim == 2:
+ audio = audio.mean(axis=1)
+ if sr != 16000:
+ import librosa
+ audio = librosa.resample(audio, orig_sr=sr, target_sr=16000)
+
+ max_samples = int(max_seconds * 16000)
+ if len(audio) > max_samples:
+ audio = audio[:max_samples]
+
+ try:
+ output = funasr_model.generate(
+ audio, granularity="utterance", extract_embedding=False,
+ )
+ except Exception:
+ return {label: 1.0 / len(PROJECT_LABELS) for label in PROJECT_LABELS}
+
+ scores = {label: 0.0 for label in PROJECT_LABELS}
+ if output and isinstance(output, list) and len(output) > 0:
+ rec = output[0]
+ raw_labels = rec.get("labels", [])
+ raw_scores = rec.get("scores", [])
+ for native_label, score in zip(raw_labels, raw_scores):
+ project_label = LABEL_MAP.get(native_label, "neutral")
+ scores[project_label] += float(score)
+
+ total = sum(scores.values())
+ if total > 0:
+ scores = {k: v / total for k, v in scores.items()}
+ return scores
+
+
+def predict_audio_onnx(audio_path: str, session, max_seconds: float = 15.0) -> dict[str, float]:
+ """Run ONNX audio emotion prediction (trimmed to max_seconds to avoid OOM)."""
+ import soundfile as sf
+
+ audio, sr = sf.read(audio_path, dtype="float32")
+ if audio.ndim == 2:
+ audio = audio.mean(axis=1)
+ if sr != 16000:
+ import librosa
+ audio = librosa.resample(audio, orig_sr=sr, target_sr=16000)
+
+ # Trim to max_seconds to prevent OOM on very long audio
+ max_samples = int(max_seconds * 16000)
+ if len(audio) > max_samples:
+ audio = audio[:max_samples]
+
+ waveform = audio.reshape(1, -1).astype(np.float32)
+ logits = session.run(None, {"waveform": waveform})[0]
+
+ exp_logits = np.exp(logits - logits.max(axis=-1, keepdims=True))
+ probs = (exp_logits / exp_logits.sum(axis=-1, keepdims=True)).squeeze()
+
+ scores = {}
+ for lora_label, prob in zip(LORA_LABELS, probs):
+ project_label = LORA_TO_PROJECT[lora_label]
+ scores[project_label] = float(prob)
+ return scores
+
+
+def predict_text_onnx(text: str, tokenizer, session) -> dict[str, float]:
+ """Run fine-tuned KcELECTRA ONNX text emotion prediction (7-class direct)."""
+ if not text or not text.strip():
+ return {label: 1.0 / len(PROJECT_LABELS) for label in PROJECT_LABELS}
+
+ enc = tokenizer(text, return_tensors="np", truncation=True, max_length=128, padding="max_length")
+ logits = session.run(None, {
+ "input_ids": enc["input_ids"],
+ "attention_mask": enc["attention_mask"],
+ })[0]
+
+ # Softmax
+ exp_logits = np.exp(logits - logits.max(axis=-1, keepdims=True))
+ probs = (exp_logits / exp_logits.sum(axis=-1, keepdims=True)).squeeze()
+
+ # LoRA KcELECTRA labels โ project labels (happiness โ joy)
+ text_labels = ["happiness", "anger", "disgust", "fear", "neutral", "sadness", "surprise"]
+ text_to_project = {
+ "happiness": "joy", "anger": "anger", "disgust": "disgust",
+ "fear": "fear", "neutral": "neutral", "sadness": "sadness", "surprise": "surprise",
+ }
+ scores = {label: 0.0 for label in PROJECT_LABELS}
+ for tl, prob in zip(text_labels, probs):
+ pl = text_to_project[tl]
+ scores[pl] = float(prob)
+ return scores
+
+
+def fuse_scores(audio_scores, text_scores, weights):
+ """Fuse with emotion-specific weights."""
+ fused = {}
+ for label in PROJECT_LABELS:
+ aw = weights.get(label, {}).get("audio", 0.6)
+ tw = weights.get(label, {}).get("text", 0.4)
+ fused[label] = audio_scores.get(label, 0.0) * aw + text_scores.get(label, 0.0) * tw
+
+ total = sum(fused.values())
+ if total > 0:
+ fused = {k: v / total for k, v in fused.items()}
+ return fused
+
+
+def compute_f1(y_true, y_pred, target_label):
+ """Compute F1 for a specific label (binary: target vs rest)."""
+ tp = sum(1 for t, p in zip(y_true, y_pred) if t == target_label and p == target_label)
+ fp = sum(1 for t, p in zip(y_true, y_pred) if t != target_label and p == target_label)
+ fn = sum(1 for t, p in zip(y_true, y_pred) if t == target_label and p != target_label)
+ precision = tp / (tp + fp) if (tp + fp) > 0 else 0
+ recall = tp / (tp + fn) if (tp + fn) > 0 else 0
+ if precision + recall == 0:
+ return 0.0
+ return 2 * precision * recall / (precision + recall)
+
+
+def compute_macro_f1(y_true, y_pred):
+ """Compute macro F1 across all 7 classes."""
+ f1s = [compute_f1(y_true, y_pred, label) for label in PROJECT_LABELS]
+ return np.mean(f1s)
+
+
+def grid_search(samples, audio_preds, text_preds):
+ """Run grid search for emotion-specific weights.
+
+ Returns:
+ grid_results: dict[emotion] โ list of {"audio_weight": float, "f1": float}
+ optimal_weights: dict[emotion] โ {"audio": float, "text": float, "f1": float}
+ """
+ weight_range = np.arange(0.0, 1.05, 0.05)
+ grid_results = {}
+ optimal_weights = {}
+
+ for target_emotion in PROJECT_LABELS:
+ results = []
+ best_f1 = -1
+ best_aw = 0.6
+
+ for aw in weight_range:
+ tw = 1.0 - aw
+ # Build per-emotion weight dict: target emotion uses (aw, tw), others use 0.6/0.4
+ weights = {}
+ for label in PROJECT_LABELS:
+ if label == target_emotion:
+ weights[label] = {"audio": float(aw), "text": float(tw)}
+ else:
+ weights[label] = {"audio": 0.6, "text": 0.4}
+
+ # Predict with these weights
+ y_true = [s["label"] for s in samples]
+ y_pred = []
+ for i, s in enumerate(samples):
+ fused = fuse_scores(audio_preds[i], text_preds[i], weights)
+ pred = max(fused, key=fused.get)
+ y_pred.append(pred)
+
+ f1 = compute_f1(y_true, y_pred, target_emotion)
+ results.append({"audio_weight": round(float(aw), 2), "f1": round(f1, 4)})
+
+ if f1 > best_f1:
+ best_f1 = f1
+ best_aw = float(aw)
+
+ grid_results[target_emotion] = results
+ optimal_weights[target_emotion] = {
+ "audio": round(best_aw, 2),
+ "text": round(1.0 - best_aw, 2),
+ "f1": round(best_f1, 4),
+ }
+ logger.info("%s: optimal audio_weight=%.2f (F1=%.4f)", target_emotion, best_aw, best_f1)
+
+ return grid_results, optimal_weights
+
+
+def compute_overall_comparison(samples, audio_preds, text_preds, optimal_weights):
+ """Compare fixed 60/40 vs optimal weights on macro F1."""
+ fixed_weights = {label: {"audio": 0.6, "text": 0.4} for label in PROJECT_LABELS}
+ y_true = [s["label"] for s in samples]
+
+ # Fixed 60/40
+ y_pred_fixed = []
+ for i in range(len(samples)):
+ fused = fuse_scores(audio_preds[i], text_preds[i], fixed_weights)
+ y_pred_fixed.append(max(fused, key=fused.get))
+ fixed_macro = compute_macro_f1(y_true, y_pred_fixed)
+ fixed_per_class = {label: compute_f1(y_true, y_pred_fixed, label) for label in PROJECT_LABELS}
+
+ # Optimal
+ opt_weight_dict = {e: {"audio": w["audio"], "text": w["text"]} for e, w in optimal_weights.items()}
+ y_pred_opt = []
+ for i in range(len(samples)):
+ fused = fuse_scores(audio_preds[i], text_preds[i], opt_weight_dict)
+ y_pred_opt.append(max(fused, key=fused.get))
+ opt_macro = compute_macro_f1(y_true, y_pred_opt)
+ opt_per_class = {label: compute_f1(y_true, y_pred_opt, label) for label in PROJECT_LABELS}
+
+ # Audio-only baseline
+ y_pred_audio = []
+ for i in range(len(samples)):
+ pred = max(audio_preds[i], key=audio_preds[i].get)
+ y_pred_audio.append(pred)
+ audio_macro = compute_macro_f1(y_true, y_pred_audio)
+ audio_per_class = {label: compute_f1(y_true, y_pred_audio, label) for label in PROJECT_LABELS}
+
+ return {
+ "audio_only": {"macro_f1": round(audio_macro, 4), "per_class": {k: round(v, 4) for k, v in audio_per_class.items()}},
+ "fixed_60_40": {"macro_f1": round(fixed_macro, 4), "per_class": {k: round(v, 4) for k, v in fixed_per_class.items()}},
+ "optimal": {"macro_f1": round(opt_macro, 4), "per_class": {k: round(v, 4) for k, v in opt_per_class.items()}},
+ }
+
+
+def plot_grid_search(grid_results, optimal_weights, output_path: Path):
+ """Plot 7 subplots: weight vs F1 per emotion."""
+ import matplotlib
+ matplotlib.use("Agg")
+ import matplotlib.pyplot as plt
+
+ fig, axes = plt.subplots(2, 4, figsize=(18, 9))
+ axes = axes.flatten()
+
+ for i, emotion in enumerate(PROJECT_LABELS):
+ ax = axes[i]
+ data = grid_results[emotion]
+ weights = [d["audio_weight"] for d in data]
+ f1s = [d["f1"] for d in data]
+ opt = optimal_weights[emotion]
+
+ ax.plot(weights, f1s, "b-o", markersize=3, linewidth=1.5)
+ ax.axvline(x=opt["audio"], color="r", linestyle="--", alpha=0.7,
+ label=f"optimal={opt['audio']:.2f}")
+ ax.axvline(x=0.6, color="gray", linestyle=":", alpha=0.5, label="fixed=0.60")
+ ax.set_title(f"{emotion} (best F1={opt['f1']:.3f})", fontsize=11, fontweight="bold")
+ ax.set_xlabel("Audio Weight")
+ ax.set_ylabel("F1 Score")
+ ax.set_xlim(-0.05, 1.05)
+ ax.legend(fontsize=8)
+ ax.grid(True, alpha=0.3)
+
+ # Hide last subplot (2x4 = 8, but only 7 emotions)
+ axes[7].set_visible(False)
+
+ fig.suptitle("Emotion-Specific Fusion Weight Grid Search", fontsize=14, fontweight="bold")
+ plt.tight_layout()
+ plt.savefig(str(output_path), dpi=150)
+ plt.close()
+ logger.info("Grid search plot saved: %s", output_path)
+
+
+def plot_comparison(comparison, optimal_weights, output_path: Path):
+ """Bar chart: audio-only vs fixed 60/40 vs optimal per emotion."""
+ import matplotlib
+ matplotlib.use("Agg")
+ import matplotlib.pyplot as plt
+
+ emotions = PROJECT_LABELS
+ audio_f1s = [comparison["audio_only"]["per_class"][e] for e in emotions]
+ fixed_f1s = [comparison["fixed_60_40"]["per_class"][e] for e in emotions]
+ opt_f1s = [comparison["optimal"]["per_class"][e] for e in emotions]
+
+ x = np.arange(len(emotions))
+ width = 0.25
+
+ fig, ax = plt.subplots(figsize=(12, 6))
+ bars1 = ax.bar(x - width, audio_f1s, width, label=f"Audio Only (macro={comparison['audio_only']['macro_f1']:.3f})", color="#2196F3", alpha=0.8)
+ bars2 = ax.bar(x, fixed_f1s, width, label=f"Fixed 60/40 (macro={comparison['fixed_60_40']['macro_f1']:.3f})", color="#FF9800", alpha=0.8)
+ bars3 = ax.bar(x + width, opt_f1s, width, label=f"Optimal (macro={comparison['optimal']['macro_f1']:.3f})", color="#4CAF50", alpha=0.8)
+
+ # Add weight annotations on optimal bars
+ for i, e in enumerate(emotions):
+ aw = optimal_weights[e]["audio"]
+ ax.text(x[i] + width, opt_f1s[i] + 0.01, f"a={aw:.0%}", ha="center", fontsize=7, color="#2E7D32")
+
+ ax.set_ylabel("F1 Score")
+ ax.set_title("Fusion Strategy Comparison: Audio Only vs Fixed 60/40 vs Emotion-Specific Optimal", fontweight="bold")
+ ax.set_xticks(x)
+ ax.set_xticklabels(emotions, fontsize=10)
+ ax.legend(fontsize=10)
+ ax.set_ylim(0, 1.0)
+ ax.grid(axis="y", alpha=0.3)
+
+ plt.tight_layout()
+ plt.savefig(str(output_path), dpi=150)
+ plt.close()
+ logger.info("Comparison plot saved: %s", output_path)
+
+
+def write_report(comparison, optimal_weights, output_path: Path):
+ """Write markdown summary report."""
+ lines = [
+ "# Fusion Weight Optimization Report",
+ "",
+ "## Summary",
+ "",
+ f"| Strategy | Macro F1 |",
+ f"|---|---|",
+ f"| Audio Only | {comparison['audio_only']['macro_f1']:.4f} |",
+ f"| Fixed 60/40 | {comparison['fixed_60_40']['macro_f1']:.4f} |",
+ f"| **Emotion-Specific Optimal** | **{comparison['optimal']['macro_f1']:.4f}** |",
+ f"| Improvement over Fixed | **+{comparison['optimal']['macro_f1'] - comparison['fixed_60_40']['macro_f1']:.4f}** |",
+ "",
+ "## Optimal Weights Per Emotion",
+ "",
+ "| Emotion | Audio Weight | Text Weight | F1 (optimal) | F1 (fixed 60/40) | Delta |",
+ "|---|---|---|---|---|---|",
+ ]
+ for e in PROJECT_LABELS:
+ aw = optimal_weights[e]["audio"]
+ tw = optimal_weights[e]["text"]
+ opt_f1 = comparison["optimal"]["per_class"][e]
+ fixed_f1 = comparison["fixed_60_40"]["per_class"][e]
+ delta = opt_f1 - fixed_f1
+ sign = "+" if delta >= 0 else ""
+ lines.append(f"| {e} | {aw:.0%} | {tw:.0%} | {opt_f1:.4f} | {fixed_f1:.4f} | {sign}{delta:.4f} |")
+
+ lines.extend([
+ "",
+ "## Methodology",
+ "",
+ "- **Data:** AI Hub 263 val split (1,294 samples, 7-class, speaker-isolated)",
+ "- **Audio model:** LoRA emotion2vec ONNX (7-class, macro F1=0.552)",
+ "- **Text model:** KcELECTRA LoRA fine-tuned (beomi/KcELECTRA-base-v2022, 7-class direct)",
+ "- **Search:** Per-emotion audio weight 0.0~1.0 in 0.05 steps (21 points ร 7 emotions)",
+ "- **Metric:** Per-emotion F1 score on val set",
+ "",
+ "## Files",
+ "",
+ "- `fusion_grid_search.json` โ full weight-F1 curve data",
+ "- `optimal_fusion_weights.json` โ best weights",
+ "- `fusion_grid_search.png` โ per-emotion weight vs F1 plots",
+ "- `fusion_comparison.png` โ strategy comparison bar chart",
+ ])
+
+ output_path.write_text("\n".join(lines), encoding="utf-8")
+ logger.info("Report saved: %s", output_path)
+
+
+def predict_text_distilroberta(text: str, tokenizer, model) -> dict[str, float]:
+ """Run j-hartmann/DistilRoBERTa text emotion prediction (7-class direct)."""
+ import torch
+
+ if not text or not text.strip():
+ return {label: 1.0 / len(PROJECT_LABELS) for label in PROJECT_LABELS}
+
+ inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
+ with torch.no_grad():
+ outputs = model(**inputs)
+ probs = torch.softmax(outputs.logits, dim=-1).squeeze().cpu().numpy()
+
+ # DistilRoBERTa labels: anger, disgust, fear, joy, neutral, sadness, surprise
+ dr_labels = [model.config.id2label[i] for i in range(len(probs))]
+ scores = {label: 0.0 for label in PROJECT_LABELS}
+ for dl, prob in zip(dr_labels, probs):
+ if dl in scores:
+ scores[dl] = float(prob)
+ return scores
+
+
+def main():
+ parser = argparse.ArgumentParser(description="Optimize emotion-specific fusion weights")
+ parser.add_argument("--lang", default="ko", choices=["ko", "en"], help="Language: ko=Korean, en=English")
+ parser.add_argument("--val-manifest", type=Path, default=Path("data/lora_dataset/val_manifest.json"))
+ parser.add_argument("--onnx-model", type=Path, default=Path("data/models/lora_emotion2vec_7class/model.onnx"))
+ parser.add_argument("--anchor-dir", type=Path, default=Path("data/AI Hub 263"))
+ parser.add_argument("--output-dir", type=Path, default=Path("data/models/fusion_optimization"))
+ parser.add_argument("--text-onnx", type=Path, default=Path("data/models/lora_kcelectra_7class/model.onnx"))
+ parser.add_argument("--text-tokenizer", default="data/models/lora_kcelectra_7class/best_model")
+ parser.add_argument("--en-text-model", default="j-hartmann/emotion-english-distilroberta-base")
+ parser.add_argument("--use-base-audio", action="store_true",
+ help="Use base (non-finetuned) emotion2vec via FunASR instead of LoRA ONNX")
+ args = parser.parse_args()
+
+ lang = args.lang
+ prefix = "en_" if lang == "en" else ""
+ output_dir = args.output_dir
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ # Step 1: Load manifest
+ with open(args.val_manifest, encoding="utf-8") as f:
+ val_all = json.load(f)
+
+ if lang == "ko":
+ # Korean: 263 val only
+ samples = [s for s in val_all if s.get("source") == "263"]
+ logger.info("Korean 263 val samples: %d", len(samples))
+ else:
+ # English: MELD test (all samples have text)
+ samples = val_all
+ logger.info("English MELD test samples: %d", len(samples))
+
+ # Map label: happiness โ joy for consistency
+ for s in samples:
+ if s["label"] == "happiness":
+ s["label"] = "joy"
+
+ # Step 2: Filter samples with text
+ matched = [s for s in samples if s.get("text", "").strip()]
+
+ # Korean fallback: load from CSV if no text in manifest
+ if not matched and lang == "ko":
+ logger.info("No text in manifest, loading from 263 CSVs...")
+ texts_map = load_263_texts(args.anchor_dir)
+ for s in samples:
+ wav_id = Path(s["path"]).stem
+ text = texts_map.get(wav_id, "")
+ if text:
+ s["text"] = text
+ matched.append(s)
+
+ logger.info("Matched audio+text: %d / %d", len(matched), len(samples))
+ if len(matched) < 50:
+ logger.error("Too few matched samples.")
+ sys.exit(1)
+
+ # Step 3: Load models
+ import onnxruntime as ort
+
+ if args.use_base_audio:
+ from funasr import AutoModel
+ logger.info("Loading base emotion2vec_plus_base via FunASR (not LoRA)...")
+ funasr_model = AutoModel(model="iic/emotion2vec_plus_base", device="cpu", hub="hf")
+ audio_predict_fn = lambda path: predict_audio_base(path, funasr_model)
+ else:
+ logger.info("Loading audio ONNX (LoRA): %s", args.onnx_model)
+ onnx_session = ort.InferenceSession(str(args.onnx_model), providers=["CPUExecutionProvider"])
+ audio_predict_fn = lambda path: predict_audio_onnx(path, onnx_session)
+
+ if lang == "ko":
+ from transformers import AutoTokenizer
+ logger.info("Loading KcELECTRA ONNX: %s", args.text_onnx)
+ text_session = ort.InferenceSession(str(args.text_onnx), providers=["CPUExecutionProvider"])
+ tokenizer = AutoTokenizer.from_pretrained(args.text_tokenizer)
+ text_predict_fn = lambda text: predict_text_onnx(text, tokenizer, text_session)
+ else:
+ from transformers import AutoTokenizer, AutoModelForSequenceClassification
+ logger.info("Loading DistilRoBERTa: %s", args.en_text_model)
+ en_tokenizer = AutoTokenizer.from_pretrained(args.en_text_model)
+ en_model = AutoModelForSequenceClassification.from_pretrained(args.en_text_model)
+ en_model.eval()
+ text_predict_fn = lambda text: predict_text_distilroberta(text, en_tokenizer, en_model)
+
+ # Step 4: Predict all samples (with checkpoint for resume safety)
+ preds_cache_path = output_dir / f"{prefix}preds_cache.json"
+ audio_preds = []
+ text_preds = []
+ start_idx = 0
+
+ if preds_cache_path.exists():
+ with open(preds_cache_path) as f:
+ cache = json.load(f)
+ audio_preds = cache.get("audio_preds", [])
+ text_preds = cache.get("text_preds", [])
+ start_idx = len(audio_preds)
+ logger.info("Resumed from checkpoint: %d predictions already done", start_idx)
+
+ for i in range(start_idx, len(matched)):
+ s = matched[i]
+ audio_scores = audio_predict_fn(s["path"])
+ audio_preds.append(audio_scores)
+
+ text_scores = text_predict_fn(s["text"])
+ text_preds.append(text_scores)
+
+ # FunASR/PyTorch leak audio tensors across .generate() calls โ force release every 25 samples
+ if (i + 1) % 25 == 0:
+ gc.collect()
+ try:
+ import torch
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+ except ImportError:
+ pass
+
+ # Checkpoint every 100 samples
+ if (i + 1) % 100 == 0:
+ logger.info("Predicted %d / %d (saving checkpoint)", i + 1, len(matched))
+ with open(preds_cache_path, "w") as f:
+ json.dump({"audio_preds": audio_preds, "text_preds": text_preds}, f)
+
+ # Final checkpoint save
+ with open(preds_cache_path, "w") as f:
+ json.dump({"audio_preds": audio_preds, "text_preds": text_preds}, f)
+
+ logger.info("All predictions done (%d samples)", len(matched))
+
+ # Step 5: Grid search
+ grid_results, optimal_weights = grid_search(matched, audio_preds, text_preds)
+
+ # Step 6: Overall comparison
+ comparison = compute_overall_comparison(matched, audio_preds, text_preds, optimal_weights)
+ logger.info("Audio-only macro F1: %.4f", comparison["audio_only"]["macro_f1"])
+ logger.info("Fixed 60/40 macro F1: %.4f", comparison["fixed_60_40"]["macro_f1"])
+ logger.info("Optimal macro F1: %.4f", comparison["optimal"]["macro_f1"])
+
+ # Step 7: Save everything
+ with open(output_dir / f"{prefix}fusion_grid_search.json", "w") as f:
+ json.dump(grid_results, f, indent=2)
+ with open(output_dir / f"{prefix}optimal_fusion_weights.json", "w") as f:
+ json.dump(optimal_weights, f, indent=2, ensure_ascii=False)
+ with open(output_dir / f"{prefix}fusion_comparison.json", "w") as f:
+ json.dump(comparison, f, indent=2)
+
+ # Step 8: Plots + report
+ plot_grid_search(grid_results, optimal_weights, output_dir / f"{prefix}fusion_grid_search.png")
+ plot_comparison(comparison, optimal_weights, output_dir / f"{prefix}fusion_comparison.png")
+ write_report(comparison, optimal_weights, output_dir / f"{prefix}fusion_report.md")
+
+ logger.info("Done! All results saved to %s", output_dir)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/prepare_aihub_test_subset.py b/scripts/prepare_aihub_test_subset.py
new file mode 100644
index 0000000000000000000000000000000000000000..236a9e77c4fcb378ce6d08553f0da2c38b4ed8eb
--- /dev/null
+++ b/scripts/prepare_aihub_test_subset.py
@@ -0,0 +1,421 @@
+#!/usr/bin/env python3
+"""AI Hub ๊ฐ์ ํ๊น
์์ ๋ํ(์ฑ์ธ) ๋ฐ์ดํฐ์
โ ๋ฒค์น๋งํฌ ํ
์คํธ ์๋ธ์
์ค๋น.
+
+AI Hub #71631 ๋ฐ์ดํฐ์
์ JSON ๋ผ๋ฒจ + ์คํ
๋ ์ค WAV์์ ๋ฐํ ๋จ์๋ฅผ ์ถ์ถํ์ฌ
+๊ท ํ ์กํ 6-class ํ
์คํธ์
์ ์์ฑํ๋ค.
+
+Usage:
+ python scripts/prepare_aihub_test_subset.py --aihub-dir data/samples
+ python scripts/prepare_aihub_test_subset.py --aihub-dir /path/to/full/dataset --samples-per-class 83
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import logging
+import os
+import random
+import sys
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Label Mapping: AI Hub ํ๊ตญ์ด โ Project 6-class
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+AIHUB_LABEL_MAP = {
+ "๊ธฐ์จ": "joy",
+ "๋๋ผ์": "surprise",
+ "๋๋ ค์": "fear",
+ "์ฌ๋์ค๋ฌ์": "joy", # Affection โ joy (user confirmed)
+ "์ฌํ": "sadness",
+ "ํ๋จ": "anger",
+ "์์": "neutral",
+ "์ค๋ฆฝ": "neutral", # appears in SpeakerEmotionCategory
+}
+
+EVAL_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear"]
+
+# Minimum utterance duration (seconds) โ too short = unreliable emotion
+MIN_DURATION_SEC = 0.5
+# Maximum utterance duration โ cap very long utterances
+MAX_DURATION_SEC = 30.0
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Step 1: Parse AI Hub JSON + discover WAV pairs
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def discover_pairs(aihub_dir: str) -> list[tuple[Path, Path]]:
+ """Find matched WAV-JSON file pairs in AI Hub directory structure.
+
+ Expected structure:
+ aihub_dir/01.์์ฒ๋ฐ์ดํฐ/{01.์ค๋ด,02.์ค์ธ}/xxx.wav
+ aihub_dir/02.๋ผ๋ฒจ๋ง๋ฐ์ดํฐ/{01.์ค๋ด,02.์ค์ธ}/xxx.json
+ """
+ source_dir = Path(aihub_dir) / "01.์์ฒ๋ฐ์ดํฐ"
+ label_dir = Path(aihub_dir) / "02.๋ผ๋ฒจ๋ง๋ฐ์ดํฐ"
+
+ if not source_dir.exists() or not label_dir.exists():
+ logger.error("Expected 01.์์ฒ๋ฐ์ดํฐ and 02.๋ผ๋ฒจ๋ง๋ฐ์ดํฐ under %s", aihub_dir)
+ sys.exit(1)
+
+ # Build WAV lookup: stem โ path
+ wav_lookup = {}
+ for wav_path in source_dir.rglob("*.wav"):
+ if str(wav_path).endswith(":Zone.Identifier"):
+ continue
+ wav_lookup[wav_path.stem] = wav_path
+
+ # Match JSON โ WAV
+ pairs = []
+ for json_path in label_dir.rglob("*.json"):
+ if str(json_path).endswith(":Zone.Identifier"):
+ continue
+ stem = json_path.stem
+ wav_path = wav_lookup.get(stem)
+ if wav_path:
+ pairs.append((wav_path, json_path))
+ else:
+ logger.warning("No WAV match for %s", json_path.name)
+
+ logger.info("Discovered %d WAV-JSON pairs", len(pairs))
+ return pairs
+
+
+def parse_utterances(pairs: list[tuple[Path, Path]]) -> list[dict]:
+ """Parse all utterances from JSON label files.
+
+ Uses VerifyEmotionTarget as ground truth (annotator-verified label).
+ """
+ utterances = []
+
+ for wav_path, json_path in pairs:
+ with open(json_path, encoding="utf-8") as f:
+ data = json.load(f)
+
+ wav_info = data.get("Wav", {})
+ file_info = data.get("File", {})
+ sr = int(wav_info.get("SamplingRate", 16000))
+ n_channels = int(wav_info.get("NumberOfChannel", 2))
+
+ # Speaker info
+ speakers = {}
+ for key in ("Speaker1", "Speaker2"):
+ spk = data.get(key, {})
+ speakers[key] = {
+ "id": spk.get("ID", ""),
+ "gender": spk.get("Gender", ""),
+ "age": spk.get("Age", ""),
+ }
+
+ for utt in data.get("Conversation", []):
+ emotion_kr = utt.get("VerifyEmotionTarget", "").strip()
+ emotion_en = AIHUB_LABEL_MAP.get(emotion_kr)
+ if emotion_en is None:
+ continue # Unknown label, skip
+
+ if emotion_en not in EVAL_LABELS:
+ continue
+
+ try:
+ start = float(str(utt["StartTime"]).replace(",", ""))
+ end = float(str(utt["EndTime"]).replace(",", ""))
+ except (KeyError, ValueError):
+ continue
+
+ duration = end - start
+ if duration < MIN_DURATION_SEC or duration > MAX_DURATION_SEC:
+ continue
+
+ speaker_no = utt.get("SpeakerNo", "Speaker1")
+ speaker_info = speakers.get(speaker_no, {})
+
+ # Determine which channel to extract (0-indexed)
+ # Speaker1 = left channel (0), Speaker2 = right channel (1)
+ channel = 0 if speaker_no == "Speaker1" else 1
+ if n_channels == 1:
+ channel = 0
+
+ utterances.append({
+ "wav_path": str(wav_path),
+ "json_path": str(json_path),
+ "file_stem": wav_path.stem,
+ "text_no": utt.get("TextNo", ""),
+ "text": utt.get("Text", ""),
+ "start": start,
+ "end": end,
+ "duration": duration,
+ "emotion": emotion_en,
+ "emotion_kr": emotion_kr,
+ "intensity": utt.get("VerifyEmotionLevel", ""),
+ "speaker_no": speaker_no,
+ "speaker_id": speaker_info.get("id", ""),
+ "speaker_gender": speaker_info.get("gender", ""),
+ "speaker_age": speaker_info.get("age", ""),
+ "channel": channel,
+ "sample_rate": sr,
+ "n_channels": n_channels,
+ })
+
+ logger.info("Parsed %d valid utterances across %d files", len(utterances), len(pairs))
+ return utterances
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Step 2: Balanced sampling
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def balanced_sample(
+ utterances: list[dict],
+ samples_per_class: int,
+ seed: int = 42,
+) -> list[dict]:
+ """Stratified balanced sampling: target samples_per_class per emotion.
+
+ Ensures:
+ - Duration diversity (short/medium/long mix)
+ - Speaker diversity (spread across speakers)
+ - For rare classes (e.g., fear), takes all available if < target
+ """
+ rng = random.Random(seed)
+
+ # Group by emotion
+ by_emotion: dict[str, list[dict]] = defaultdict(list)
+ for utt in utterances:
+ by_emotion[utt["emotion"]].append(utt)
+
+ selected = []
+ stats = {}
+
+ for emotion in EVAL_LABELS:
+ pool = by_emotion.get(emotion, [])
+ if not pool:
+ logger.warning("No samples for emotion '%s'", emotion)
+ stats[emotion] = 0
+ continue
+
+ if len(pool) <= samples_per_class:
+ # Take all for rare classes
+ chosen = pool
+ else:
+ # Duration-stratified sampling
+ short = [u for u in pool if u["duration"] < 3.0]
+ medium = [u for u in pool if 3.0 <= u["duration"] < 10.0]
+ long = [u for u in pool if u["duration"] >= 10.0]
+
+ # Target ratio: 30% short, 50% medium, 20% long
+ n_short = max(1, int(samples_per_class * 0.3))
+ n_long = max(1, int(samples_per_class * 0.2))
+ n_medium = samples_per_class - n_short - n_long
+
+ chosen = []
+ for bucket, n in [(short, n_short), (medium, n_medium), (long, n_long)]:
+ rng.shuffle(bucket)
+ chosen.extend(bucket[:n])
+
+ # Fill remaining if any bucket was short
+ if len(chosen) < samples_per_class:
+ remaining = [u for u in pool if u not in chosen]
+ rng.shuffle(remaining)
+ chosen.extend(remaining[: samples_per_class - len(chosen)])
+
+ chosen = chosen[:samples_per_class]
+
+ selected.extend(chosen)
+ stats[emotion] = len(chosen)
+
+ logger.info("Sampling result: %s (total: %d)", stats, len(selected))
+ return selected
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Step 3: Extract utterance WAVs
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def extract_utterances(
+ selected: list[dict],
+ output_dir: str,
+) -> list[dict]:
+ """Extract individual utterance WAV segments from conversation files.
+
+ Reads the stereo WAV, extracts the correct speaker channel,
+ and saves as mono 16kHz WAV.
+ """
+ out_path = Path(output_dir)
+ records = []
+
+ # Cache loaded audio files (avoid re-reading same WAV)
+ audio_cache: dict[str, tuple[np.ndarray, int]] = {}
+
+ for i, utt in enumerate(selected):
+ emotion = utt["emotion"]
+ emotion_dir = out_path / "test_audio" / emotion
+ emotion_dir.mkdir(parents=True, exist_ok=True)
+
+ # Load audio (cached)
+ wav_path = utt["wav_path"]
+ if wav_path not in audio_cache:
+ try:
+ audio, sr = sf.read(wav_path, dtype="float32")
+ audio_cache[wav_path] = (audio, sr)
+ except Exception as e:
+ logger.warning("Failed to read %s: %s", wav_path, e)
+ continue
+
+ audio, sr = audio_cache[wav_path]
+
+ # Extract channel
+ if audio.ndim == 2:
+ channel = min(utt["channel"], audio.shape[1] - 1)
+ mono = audio[:, channel]
+ else:
+ mono = audio
+
+ # Extract time range
+ start_sample = int(utt["start"] * sr)
+ end_sample = int(utt["end"] * sr)
+ start_sample = max(0, start_sample)
+ end_sample = min(len(mono), end_sample)
+
+ segment = mono[start_sample:end_sample]
+
+ if len(segment) < int(MIN_DURATION_SEC * sr):
+ logger.warning("Segment too short after extraction: %s_%s", utt["file_stem"], utt["text_no"])
+ continue
+
+ # Resample to 16kHz if needed
+ if sr != 16000:
+ import librosa
+ segment = librosa.resample(segment, orig_sr=sr, target_sr=16000)
+ sr = 16000
+
+ # Save
+ filename = f"kr_{emotion}_{i:04d}.wav"
+ filepath = emotion_dir / filename
+ sf.write(str(filepath), segment, 16000, subtype="PCM_16")
+
+ records.append({
+ "file_path": str(filepath.relative_to(out_path)),
+ "emotion": emotion,
+ "duration": round(len(segment) / 16000, 3),
+ "speaker_id": utt["speaker_id"],
+ "speaker_gender": utt["speaker_gender"],
+ "intensity": utt["intensity"],
+ "text": utt["text"],
+ "source_file": utt["file_stem"],
+ })
+
+ if (i + 1) % 100 == 0:
+ logger.info("Extracted %d/%d utterances", i + 1, len(selected))
+
+ logger.info("Extracted %d utterance WAVs to %s", len(records), output_dir)
+ return records
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Step 4: Write labels CSV + metadata JSON
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def write_outputs(records: list[dict], output_dir: str, utterances: list[dict]):
+ """Write test_labels.csv and metadata.json."""
+ out_path = Path(output_dir)
+
+ # CSV
+ csv_path = out_path / "test_labels.csv"
+ with open(csv_path, "w", newline="", encoding="utf-8") as f:
+ writer = csv.DictWriter(f, fieldnames=[
+ "file_path", "emotion", "duration", "speaker_id",
+ "speaker_gender", "intensity", "text", "source_file",
+ ])
+ writer.writeheader()
+ writer.writerows(records)
+ logger.info("Wrote %s (%d records)", csv_path, len(records))
+
+ # Metadata
+ from collections import Counter
+ emotion_dist = Counter(r["emotion"] for r in records)
+ duration_stats = [r["duration"] for r in records]
+ intensity_dist = Counter(r["intensity"] for r in records)
+
+ metadata = {
+ "dataset": "AI Hub #71631 - ๊ฐ์ ์ด ํ๊น
๋ ์์ ๋ํ (์ฑ์ธ)",
+ "subset": "test",
+ "total_samples": len(records),
+ "eval_classes": EVAL_LABELS,
+ "label_mapping": AIHUB_LABEL_MAP,
+ "emotion_distribution": dict(emotion_dist),
+ "intensity_distribution": dict(intensity_dist),
+ "duration_stats": {
+ "mean": round(sum(duration_stats) / max(len(duration_stats), 1), 2),
+ "min": round(min(duration_stats, default=0), 2),
+ "max": round(max(duration_stats, default=0), 2),
+ },
+ "total_source_utterances": len(utterances),
+ "note": "disgust class absent from AI Hub dataset โ 6-class evaluation",
+ }
+
+ meta_path = out_path / "metadata.json"
+ with open(meta_path, "w", encoding="utf-8") as f:
+ json.dump(metadata, f, indent=2, ensure_ascii=False)
+ logger.info("Wrote %s", meta_path)
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Main
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def main():
+ parser = argparse.ArgumentParser(
+ description="AI Hub ๊ฐ์ ๋ฐ์ดํฐ์
โ ๋ฒค์น๋งํฌ ํ
์คํธ ์๋ธ์
์ค๋น",
+ )
+ parser.add_argument("--aihub-dir", required=True, help="AI Hub ๋ฐ์ดํฐ ๋ฃจํธ (01.์์ฒ๋ฐ์ดํฐ, 02.๋ผ๋ฒจ๋ง๋ฐ์ดํฐ ํฌํจ)")
+ parser.add_argument("--output-dir", default="data/evaluation/korean", help="์ถ๋ ฅ ๋๋ ํ ๋ฆฌ")
+ parser.add_argument("--samples-per-class", type=int, default=83, help="ํด๋์ค๋น ๋ชฉํ ์ํ ์ (default: 83)")
+ parser.add_argument("--seed", type=int, default=42, help="๋๋ค ์๋")
+ parser.add_argument("--ground-truth", default="verify", choices=["verify", "speaker"],
+ help="Ground truth ์์ค: verify=๊ฒ์ฆ์ ๋ผ๋ฒจ, speaker=ํ์ ์๊ธฐ๋ณด๊ณ ")
+ args = parser.parse_args()
+
+ # 1. Discover pairs
+ pairs = discover_pairs(args.aihub_dir)
+ if not pairs:
+ logger.error("No WAV-JSON pairs found")
+ sys.exit(1)
+
+ # 2. Parse utterances
+ utterances = parse_utterances(pairs)
+ if not utterances:
+ logger.error("No valid utterances parsed")
+ sys.exit(1)
+
+ # Log distribution before sampling
+ from collections import Counter
+ raw_dist = Counter(u["emotion"] for u in utterances)
+ logger.info("Raw distribution: %s", dict(raw_dist))
+
+ # 3. Balanced sampling
+ selected = balanced_sample(utterances, args.samples_per_class, seed=args.seed)
+
+ # 4. Extract WAVs
+ records = extract_utterances(selected, args.output_dir)
+
+ # 5. Write outputs
+ write_outputs(records, args.output_dir, utterances)
+
+ print(f"\nDone! Test subset ready at {args.output_dir}/")
+ print(f" - {len(records)} utterance WAVs in test_audio/")
+ print(f" - test_labels.csv")
+ print(f" - metadata.json")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/prepare_dataset.py b/scripts/prepare_dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/scripts/prepare_lora_dataset.py b/scripts/prepare_lora_dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..6e413dded63d95e7a62fe2f0ec49e85f7cc3e2fe
--- /dev/null
+++ b/scripts/prepare_lora_dataset.py
@@ -0,0 +1,875 @@
+#!/usr/bin/env python3
+"""Unified data preprocessing pipeline for LoRA emotion2vec 7-class fine-tuning.
+
+Extracts, preprocesses, and merges samples from three sources:
+ 1. AI Hub 263 โ anchor (acted Korean, 7-class)
+ 2. AI Hub 71631 โ booster (outdoor spontaneous Korean, mapped to 7-class)
+ 3. RAVDESS โ English (acted, 7-class)
+
+Outputs a unified manifest (train/val JSONs) ready for LoRA training.
+
+Usage:
+ python scripts/prepare_lora_dataset.py \
+ --anchor-dir "data/AI Hub 263" \
+ --booster-label-zip "data/AI Hub 71631/01-1.์ ์๊ฐ๋ฐฉ๋ฐ์ดํฐ/Training/02.๋ผ๋ฒจ๋ง๋ฐ์ดํฐ/TL_02.์ค์ธ.zip" \
+ --booster-audio-zip "data/AI Hub 71631/01-1.์ ์๊ฐ๋ฐฉ๋ฐ์ดํฐ/Training/01.์์ฒ๋ฐ์ดํฐ/TS_02.์ค์ธ.zip" \
+ --ravdess-dir data/ravdess \
+ --output-dir data/lora_7class
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import io
+import json
+import logging
+import random
+import zipfile
+from collections import Counter, defaultdict
+from pathlib import Path
+
+import torch
+import torchaudio
+
+logging.basicConfig(
+ level=logging.INFO,
+ format="%(asctime)s %(levelname)-8s %(message)s",
+ datefmt="%H:%M:%S",
+)
+log = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+LABEL2IDX: dict[str, int] = {
+ "happiness": 0,
+ "anger": 1,
+ "disgust": 2,
+ "fear": 3,
+ "neutral": 4,
+ "sadness": 5,
+ "surprise": 6,
+}
+
+VALID_LABELS = set(LABEL2IDX.keys())
+
+TARGET_SR = 16_000
+RMS_THRESHOLD = 0.001 # 0.005โ0.001: disgust ๋ฑ ์ ์๋์ง ๋ฐํ ๋ณด์กด (์ง์ง ๋ฌด์๋ง ์ ๊ฑฐ)
+
+# ---------------------------------------------------------------------------
+# Task 1 โ Label Mappers
+# ---------------------------------------------------------------------------
+
+_MAP_263: dict[str, str] = {
+ "angry": "anger",
+ "happiness": "happiness",
+ "neutral": "neutral",
+ "sadness": "sadness",
+ "surprise": "surprise",
+ "fear": "fear",
+ "disgust": "disgust",
+}
+
+
+def map_263_label(raw: str) -> str | None:
+ """Map AI Hub 263 annotator label to 7-class. Case-insensitive."""
+ if not raw:
+ return None
+ return _MAP_263.get(raw.strip().lower())
+
+
+_MAP_71631: dict[str, str] = {
+ "๊ธฐ์จ": "happiness",
+ "ํ๋จ": "anger",
+ "๋๋ผ์": "surprise",
+ "์ฌํ": "sadness",
+ "๋๋ ค์": "fear",
+ "์์": "neutral",
+ "์ค๋ฆฝ": "neutral",
+}
+
+
+def map_71631_label(raw: str) -> str | None:
+ """Map AI Hub 71631 Korean emotion label to 7-class."""
+ if not raw:
+ return None
+ return _MAP_71631.get(raw.strip())
+
+
+def map_ravdess_label(raw: str) -> str | None:
+ """Map RAVDESS label to 7-class. Passthrough except joyโhappiness."""
+ if not raw:
+ return None
+ lbl = raw.strip().lower()
+ if lbl == "joy":
+ return "happiness"
+ if lbl in VALID_LABELS:
+ return lbl
+ return None
+
+
+def majority_vote_263(emotions: list[str | None]) -> str | None:
+ """Return majority label (3/5+) or None for no majority/tie."""
+ valid = [e for e in emotions if e is not None]
+ if not valid:
+ return None
+ counts = Counter(valid)
+ top_label, top_count = counts.most_common(1)[0]
+ if top_count < 3:
+ return None
+ # Check for tie at top count
+ tied = [lbl for lbl, c in counts.items() if c == top_count]
+ if len(tied) > 1:
+ return None
+ return top_label
+
+
+# ---------------------------------------------------------------------------
+# Task 1 โ Audio Preprocessing
+# ---------------------------------------------------------------------------
+
+
+import re
+
+def _clean_text(text: str) -> str:
+ """Clean text for STT-friendly format.
+
+ Removes non-verbal tags like (์์), (ํ์จ), keeps ?, !, ...
+ """
+ # Remove non-verbal tags: (์์), (ํ์จ), (์นจ๋ฌต), [noise], etc.
+ text = re.sub(r"[(\[๏ผ][^)\]๏ผ]*[)\]๏ผ]", "", text)
+ # Remove trailing/leading whitespace, collapse multiple spaces
+ text = re.sub(r"\s+", " ", text).strip()
+ return text
+
+
+def _compute_rms(waveform: torch.Tensor) -> float:
+ """Compute RMS of a waveform tensor."""
+ return float(torch.sqrt(torch.mean(waveform.float() ** 2)))
+
+
+def preprocess_audio(input_path: Path, output_path: Path) -> bool:
+ """Resample to 16kHz mono, trim silence, reject if RMS < threshold.
+
+ Returns True if file was saved, False if rejected.
+ """
+ input_path = Path(input_path)
+ output_path = Path(output_path)
+
+ waveform, sr = torchaudio.load(str(input_path))
+
+ # Mono
+ if waveform.shape[0] > 1:
+ waveform = waveform.mean(dim=0, keepdim=True)
+
+ # Resample
+ if sr != TARGET_SR:
+ waveform = torchaudio.functional.resample(waveform, sr, TARGET_SR)
+
+ # Trim silence (leading/trailing)
+ waveform_trimmed = torchaudio.functional.vad(waveform, TARGET_SR)
+ if waveform_trimmed.numel() > 0:
+ waveform = waveform_trimmed
+
+ # RMS check
+ rms = _compute_rms(waveform)
+ if rms < RMS_THRESHOLD:
+ return False
+
+ output_path.parent.mkdir(parents=True, exist_ok=True)
+ torchaudio.save(str(output_path), waveform, TARGET_SR)
+ return True
+
+
+def preprocess_audio_from_tensor(
+ waveform: torch.Tensor, sr: int, output_path: Path
+) -> bool:
+ """Preprocess an in-memory waveform tensor and save to output_path.
+
+ Used for 71631 where we slice in memory from the full conversation wav.
+ """
+ output_path = Path(output_path)
+
+ # Mono
+ if waveform.dim() > 1 and waveform.shape[0] > 1:
+ waveform = waveform.mean(dim=0, keepdim=True)
+ elif waveform.dim() == 1:
+ waveform = waveform.unsqueeze(0)
+
+ # Resample
+ if sr != TARGET_SR:
+ waveform = torchaudio.functional.resample(waveform, sr, TARGET_SR)
+
+ # Trim silence (leading/trailing) โ same as preprocess_audio
+ waveform_trimmed = torchaudio.functional.vad(waveform, TARGET_SR)
+ if waveform_trimmed.numel() > 0:
+ waveform = waveform_trimmed
+
+ # RMS check
+ rms = _compute_rms(waveform)
+ if rms < RMS_THRESHOLD:
+ return False
+
+ output_path.parent.mkdir(parents=True, exist_ok=True)
+ torchaudio.save(str(output_path), waveform, TARGET_SR)
+ return True
+
+
+# ---------------------------------------------------------------------------
+# Task 2 โ AI Hub 263 Extraction
+# ---------------------------------------------------------------------------
+
+
+def parse_263_row(row: list[str]) -> dict | None:
+ """Parse a CSV row from AI Hub 263.
+
+ Columns: wav_id, ๋ฐํ๋ฌธ, ์ํฉ, 1๋ฒ๊ฐ์ , 1๋ฒ๊ฐ์ ์ธ๊ธฐ, 2๋ฒ๊ฐ์ , 2๋ฒ๊ฐ์ ์ธ๊ธฐ,
+ 3๋ฒ๊ฐ์ , 3๋ฒ๊ฐ์ ์ธ๊ธฐ, 4๋ฒ๊ฐ์ , 4๋ฒ๊ฐ์ ์ธ๊ธฐ, 5๋ฒ๊ฐ์ , 5๋ฒ๊ฐ์ ์ธ๊ธฐ, ๋์ด, ์ฑ๋ณ
+
+ Returns dict with wav_id, label, agreement, max_intensity, or None if no majority.
+ """
+ if len(row) < 15:
+ return None
+
+ wav_id = row[0].strip()
+ # Extract 5 annotator emotions and intensities
+ annotator_emotions: list[str | None] = []
+ intensities: list[int] = []
+ for i in range(5):
+ emo_col = 3 + i * 2 # 3, 5, 7, 9, 11
+ int_col = 3 + i * 2 + 1 # 4, 6, 8, 10, 12
+ raw_emo = row[emo_col].strip() if emo_col < len(row) else ""
+ raw_int = row[int_col].strip() if int_col < len(row) else "0"
+ mapped = map_263_label(raw_emo)
+ annotator_emotions.append(mapped)
+ try:
+ intensities.append(int(raw_int))
+ except ValueError:
+ intensities.append(0)
+
+ label = majority_vote_263(annotator_emotions)
+ if label is None:
+ return None
+
+ agreement = sum(1 for e in annotator_emotions if e == label)
+ # Max intensity among annotators who voted for the majority label
+ max_intensity = max(
+ (intensities[i] for i, e in enumerate(annotator_emotions) if e == label),
+ default=0,
+ )
+
+ # ๋ฐํ๋ฌธ (text) โ column 1
+ text = row[1].strip() if len(row) > 1 else ""
+ text = _clean_text(text)
+
+ return {
+ "wav_id": wav_id,
+ "text": text,
+ "label": label,
+ "agreement": agreement,
+ "max_intensity": max_intensity,
+ }
+
+
+def extract_anchor_263(
+ anchor_dir: Path,
+ output_dir: Path,
+ cap_per_class: int = 2100,
+) -> list[dict]:
+ """Extract and preprocess AI Hub 263 dataset.
+
+ Parses 3 CSVs (cp949), majority-votes annotator labels,
+ priority-sorts (agreement desc, intensity desc), caps per class,
+ extracts wavs from ZIPs, preprocesses audio.
+ """
+ anchor_dir = Path(anchor_dir)
+ output_dir = Path(output_dir)
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ csv_files = sorted(anchor_dir.glob("*.csv"))
+ zip_files = sorted(anchor_dir.glob("*.zip"))
+
+ log.info("263: Found %d CSVs, %d ZIPs", len(csv_files), len(zip_files))
+
+ # Step 1: Parse all CSVs
+ all_parsed: list[dict] = []
+ for csv_path in csv_files:
+ with open(csv_path, encoding="cp949", newline="") as f:
+ reader = csv.reader(f)
+ header = next(reader) # skip header
+ for row in reader:
+ result = parse_263_row(row)
+ if result is not None:
+ result["csv_source"] = csv_path.stem
+ all_parsed.append(result)
+
+ log.info("263: Parsed %d rows with majority vote", len(all_parsed))
+
+ # Step 2: Group by label, priority sort, cap
+ by_label: dict[str, list[dict]] = defaultdict(list)
+ for item in all_parsed:
+ by_label[item["label"]].append(item)
+
+ selected: list[dict] = []
+ for label, items in by_label.items():
+ # Sort by agreement desc, then intensity desc
+ items.sort(key=lambda x: (x["agreement"], x["max_intensity"]), reverse=True)
+ capped = items[:cap_per_class]
+ selected.extend(capped)
+ log.info("263: %s โ %d available, %d selected", label, len(items), len(capped))
+
+ # Step 3: Build wav_id โ zip lookup
+ wav_to_zip: dict[str, tuple[zipfile.ZipFile, str]] = {}
+ zip_handles = [zipfile.ZipFile(zp) for zp in zip_files]
+ for zf in zip_handles:
+ for name in zf.namelist():
+ if name.endswith(".wav"):
+ basename = Path(name).stem
+ wav_to_zip[basename] = (zf, name)
+
+ # Step 4: Extract and preprocess
+ samples: list[dict] = []
+ skipped = 0
+ for item in selected:
+ wav_id = item["wav_id"]
+ if wav_id not in wav_to_zip:
+ skipped += 1
+ continue
+
+ zf, zip_entry = wav_to_zip[wav_id]
+ out_path = output_dir / item["label"] / f"{wav_id}.wav"
+
+ try:
+ with zf.open(zip_entry) as src:
+ audio_bytes = src.read()
+
+ # Write to temp, then preprocess
+ tmp_path = output_dir / f"_tmp_{wav_id}.wav"
+ tmp_path.write_bytes(audio_bytes)
+
+ ok = preprocess_audio(tmp_path, out_path)
+ tmp_path.unlink(missing_ok=True)
+
+ if ok:
+ samples.append({
+ "path": str(out_path),
+ "label": item["label"],
+ "label_idx": LABEL2IDX[item["label"]],
+ "source": "263",
+ "speaker_id": f"263_{wav_id[:8]}",
+ "text": item.get("text", ""),
+ "agreement": item["agreement"],
+ "intensity": item["max_intensity"],
+ })
+ else:
+ skipped += 1
+ except Exception as e:
+ log.warning("263: Failed to process %s: %s", wav_id, e)
+ skipped += 1
+
+ # Close zip handles
+ for zf in zip_handles:
+ zf.close()
+
+ log.info("263: Extracted %d samples, skipped %d", len(samples), skipped)
+ return samples
+
+
+# ---------------------------------------------------------------------------
+# Task 3 โ AI Hub 71631 Outdoor Extraction
+# ---------------------------------------------------------------------------
+
+
+def parse_71631_utterance(conv_entry: dict) -> dict | None:
+ """Parse a conversation entry from 71631 JSON.
+
+ Filters by VerifyEmotionLevel (๋ณดํต/๊ฐํจ only, rejects ์ฝํจ).
+ Returns dict with label, intensity, start_time, end_time, speaker_no or None.
+ """
+ level = conv_entry.get("VerifyEmotionLevel", "")
+ if level not in ("๋ณดํต", "๊ฐํจ"):
+ return None
+
+ emotion = conv_entry.get("VerifyEmotionTarget", "")
+ label = map_71631_label(emotion)
+ if label is None:
+ return None
+
+ try:
+ start_time = float(conv_entry["StartTime"])
+ end_time = float(conv_entry["EndTime"])
+ except (KeyError, ValueError):
+ return None
+
+ if end_time <= start_time:
+ return None
+
+ # Text from conversation entry
+ text = _clean_text(conv_entry.get("Text", ""))
+
+ return {
+ "label": label,
+ "intensity": level,
+ "start_time": start_time,
+ "end_time": end_time,
+ "speaker_no": conv_entry.get("SpeakerNo", ""),
+ "text": text,
+ }
+
+
+def extract_booster_71631(
+ label_zip: Path,
+ audio_zip: Path,
+ output_dir: Path,
+ cap: int = 3500,
+ max_per_speaker: int = 20,
+) -> list[dict]:
+ """Extract and preprocess AI Hub 71631 outdoor dataset.
+
+ Parses label ZIP JSONs, filters intensity, applies speaker cap,
+ slices wav segments, resamples to 16kHz, RMS-filters neutral.
+ """
+ label_zip = Path(label_zip)
+ audio_zip = Path(audio_zip)
+ output_dir = Path(output_dir)
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ # Step 1: Parse all label JSONs
+ all_utterances: list[dict] = []
+ with zipfile.ZipFile(label_zip) as lzf:
+ json_files = [n for n in lzf.namelist() if n.endswith(".json")]
+ log.info("71631: Found %d label JSONs", len(json_files))
+
+ for jf in json_files:
+ try:
+ with lzf.open(jf) as f:
+ data = json.load(f)
+ except Exception as e:
+ log.warning("71631: Failed to parse %s: %s", jf, e)
+ continue
+
+ filename = data.get("File", {}).get("FileName", "")
+ conv_id = filename # Use filename as conversation_id
+ spk1_id = data.get("Speaker1", {}).get("ID", "")
+ spk2_id = data.get("Speaker2", {}).get("ID", "")
+
+ for entry in data.get("Conversation", []):
+ parsed = parse_71631_utterance(entry)
+ if parsed is None:
+ continue
+ # Determine speaker ID
+ spk_no = parsed["speaker_no"]
+ if spk_no == "Speaker1":
+ spk_id = spk1_id
+ elif spk_no == "Speaker2":
+ spk_id = spk2_id
+ else:
+ spk_id = spk_no
+
+ parsed["conversation_id"] = conv_id
+ parsed["speaker_id"] = f"71631_{spk_id}"
+ parsed["wav_filename"] = filename
+ parsed["text_no"] = entry.get("TextNo", "")
+ all_utterances.append(parsed)
+
+ log.info("71631: Parsed %d utterances (๋ณดํต/๊ฐํจ)", len(all_utterances))
+
+ # Step 2: Speaker cap
+ speaker_counts: Counter = Counter()
+ speaker_capped: list[dict] = []
+ # Priority: ๊ฐํจ first, then ๋ณดํต
+ all_utterances.sort(key=lambda x: (0 if x["intensity"] == "๊ฐํจ" else 1))
+ for utt in all_utterances:
+ spk = utt["speaker_id"]
+ if speaker_counts[spk] < max_per_speaker:
+ speaker_capped.append(utt)
+ speaker_counts[spk] += 1
+
+ log.info("71631: After speaker cap (%d/spk): %d utterances", max_per_speaker, len(speaker_capped))
+
+ # Step 3: Group by label, cap per class
+ by_label: dict[str, list[dict]] = defaultdict(list)
+ for utt in speaker_capped:
+ by_label[utt["label"]].append(utt)
+
+ selected: list[dict] = []
+ for label, items in by_label.items():
+ # Priority: ๊ฐํจ first (already sorted)
+ capped = items[:cap]
+ selected.extend(capped)
+ log.info("71631: %s โ %d available, %d selected", label, len(items), len(capped))
+
+ # Step 4: Group by wav filename for efficient audio loading
+ by_wav: dict[str, list[dict]] = defaultdict(list)
+ for utt in selected:
+ by_wav[utt["wav_filename"]].append(utt)
+
+ # Step 5: Extract audio segments
+ samples: list[dict] = []
+ skipped = 0
+
+ with zipfile.ZipFile(audio_zip) as azf:
+ wav_lookup: dict[str, str] = {}
+ for name in azf.namelist():
+ if name.endswith(".wav"):
+ stem = Path(name).stem
+ wav_lookup[stem] = name
+
+ for wav_filename, utterances in by_wav.items():
+ if wav_filename not in wav_lookup:
+ log.warning("71631: WAV not found in zip: %s", wav_filename)
+ skipped += len(utterances)
+ continue
+
+ zip_entry = wav_lookup[wav_filename]
+ try:
+ with azf.open(zip_entry) as src:
+ audio_bytes = src.read()
+
+ buf = io.BytesIO(audio_bytes)
+ waveform, sr = torchaudio.load(buf)
+ except Exception as e:
+ log.warning("71631: Failed to load %s: %s", wav_filename, e)
+ skipped += len(utterances)
+ continue
+
+ # Mono
+ if waveform.shape[0] > 1:
+ waveform = waveform.mean(dim=0, keepdim=True)
+
+ for utt in utterances:
+ start_sample = int(utt["start_time"] * sr)
+ end_sample = int(utt["end_time"] * sr)
+
+ if end_sample > waveform.shape[1]:
+ end_sample = waveform.shape[1]
+ if start_sample >= end_sample:
+ skipped += 1
+ continue
+
+ segment = waveform[:, start_sample:end_sample]
+
+ out_name = f"{wav_filename}_{utt['text_no']}.wav"
+ out_path = output_dir / utt["label"] / out_name
+
+ ok = preprocess_audio_from_tensor(segment, sr, out_path)
+ if ok:
+ samples.append({
+ "path": str(out_path),
+ "label": utt["label"],
+ "label_idx": LABEL2IDX[utt["label"]],
+ "source": "71631",
+ "speaker_id": utt["speaker_id"],
+ "text": utt.get("text", ""),
+ "conversation_id": utt["conversation_id"],
+ "intensity": utt["intensity"],
+ })
+ else:
+ skipped += 1
+
+ log.info("71631: Extracted %d samples, skipped %d", len(samples), skipped)
+ return samples
+
+
+# ---------------------------------------------------------------------------
+# Task 4 โ RAVDESS Extraction
+# ---------------------------------------------------------------------------
+
+
+def extract_ravdess(ravdess_dir: Path, output_dir: Path) -> list[dict]:
+ """Extract and preprocess RAVDESS dataset from manifest.csv."""
+ ravdess_dir = Path(ravdess_dir)
+ output_dir = Path(output_dir)
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ manifest_path = ravdess_dir / "manifest.csv"
+ if not manifest_path.exists():
+ log.error("RAVDESS manifest not found: %s", manifest_path)
+ return []
+
+ samples: list[dict] = []
+ skipped = 0
+
+ with open(manifest_path) as f:
+ reader = csv.DictReader(f)
+ for row in reader:
+ clean_path = Path(row["clean_path"])
+ emotion_raw = row.get("emotion", "")
+ label = map_ravdess_label(emotion_raw)
+ if label is None:
+ skipped += 1
+ continue
+
+ actor_id = int(row["actor_id"])
+ out_name = clean_path.name
+ out_path = output_dir / label / out_name
+
+ if not clean_path.exists():
+ skipped += 1
+ continue
+
+ ok = preprocess_audio(clean_path, out_path)
+ if ok:
+ samples.append({
+ "path": str(out_path),
+ "label": label,
+ "label_idx": LABEL2IDX[label],
+ "source": "ravdess",
+ "actor_id": actor_id,
+ "speaker_id": f"ravdess_{actor_id}",
+ "text": "", # RAVDESS uses fixed sentences, not useful for text emotion
+ })
+ else:
+ skipped += 1
+
+ log.info("RAVDESS: Extracted %d samples, skipped %d", len(samples), skipped)
+ return samples
+
+
+# ---------------------------------------------------------------------------
+# Task 4 โ Train/Val Splits
+# ---------------------------------------------------------------------------
+
+
+def speaker_isolated_split(
+ samples: list[dict], val_ratio: float = 0.1
+) -> tuple[list[dict], list[dict]]:
+ """Split by speaker_id โ no leakage between train/val."""
+ if not samples:
+ return [], []
+
+ # Group by speaker
+ by_speaker: dict[str, list[dict]] = defaultdict(list)
+ for s in samples:
+ by_speaker[s["speaker_id"]].append(s)
+
+ speakers = list(by_speaker.keys())
+ random.shuffle(speakers)
+
+ total = len(samples)
+ target_val = int(total * val_ratio)
+
+ val_samples: list[dict] = []
+ val_speakers: set[str] = set()
+ for spk in speakers:
+ if len(val_samples) >= target_val:
+ break
+ val_samples.extend(by_speaker[spk])
+ val_speakers.add(spk)
+
+ train_samples = [s for s in samples if s["speaker_id"] not in val_speakers]
+ return train_samples, val_samples
+
+
+def conversation_isolated_split(
+ samples: list[dict], val_ratio: float = 0.1
+) -> tuple[list[dict], list[dict]]:
+ """Split by conversation_id โ no leakage between train/val."""
+ if not samples:
+ return [], []
+
+ by_conv: dict[str, list[dict]] = defaultdict(list)
+ for s in samples:
+ by_conv[s["conversation_id"]].append(s)
+
+ convs = list(by_conv.keys())
+ random.shuffle(convs)
+
+ total = len(samples)
+ target_val = int(total * val_ratio)
+
+ val_samples: list[dict] = []
+ val_convs: set[str] = set()
+ for conv in convs:
+ if len(val_samples) >= target_val:
+ break
+ val_samples.extend(by_conv[conv])
+ val_convs.add(conv)
+
+ train_samples = [s for s in samples if s["conversation_id"] not in val_convs]
+ return train_samples, val_samples
+
+
+def actor_split_ravdess(
+ samples: list[dict], val_actors: list[int]
+) -> tuple[list[dict], list[dict]]:
+ """Split RAVDESS by actor โ specified actors go to val."""
+ val_set = set(val_actors)
+ train = [s for s in samples if s["actor_id"] not in val_set]
+ val = [s for s in samples if s["actor_id"] in val_set]
+ return train, val
+
+
+# ---------------------------------------------------------------------------
+# Task 4 โ Manifest & Stats
+# ---------------------------------------------------------------------------
+
+
+def save_manifest(samples: list[dict], path: Path) -> None:
+ """Save manifest as JSON Lines."""
+ path = Path(path)
+ path.parent.mkdir(parents=True, exist_ok=True)
+ with open(path, "w") as f:
+ json.dump(samples, f, indent=2, ensure_ascii=False)
+ log.info("Saved manifest: %s (%d samples)", path, len(samples))
+
+
+def save_stats(train: list[dict], val: list[dict], path: Path) -> None:
+ """Save dataset statistics."""
+ path = Path(path)
+ path.parent.mkdir(parents=True, exist_ok=True)
+
+ def _count_stats(samples: list[dict]) -> dict:
+ by_label: Counter = Counter()
+ by_source: Counter = Counter()
+ for s in samples:
+ by_label[s["label"]] += 1
+ by_source[s["source"]] += 1
+ return {
+ "total": len(samples),
+ "by_label": dict(sorted(by_label.items())),
+ "by_source": dict(sorted(by_source.items())),
+ }
+
+ stats = {
+ "train": _count_stats(train),
+ "val": _count_stats(val),
+ "label2idx": LABEL2IDX,
+ }
+
+ with open(path, "w") as f:
+ json.dump(stats, f, indent=2, ensure_ascii=False)
+ log.info("Saved stats: %s", path)
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+
+def main():
+ parser = argparse.ArgumentParser(
+ description="Prepare unified LoRA 7-class dataset"
+ )
+ parser.add_argument(
+ "--anchor-dir",
+ type=Path,
+ default=Path("data/AI Hub 263"),
+ help="Path to AI Hub 263 directory with CSVs + ZIPs",
+ )
+ parser.add_argument(
+ "--booster-label-zip",
+ type=Path,
+ default=Path(
+ "data/AI Hub 71631/01-1.์ ์๊ฐ๋ฐฉ๋ฐ์ดํฐ/Training/"
+ "02.๋ผ๋ฒจ๋ง๋ฐ์ดํฐ/TL_02.์ค์ธ.zip"
+ ),
+ help="Path to 71631 label ZIP",
+ )
+ parser.add_argument(
+ "--booster-audio-zip",
+ type=Path,
+ default=Path(
+ "data/AI Hub 71631/01-1.์ ์๊ฐ๋ฐฉ๋ฐ์ดํฐ/Training/"
+ "01.์์ฒ๋ฐ์ดํฐ/TS_02.์ค์ธ.zip"
+ ),
+ help="Path to 71631 audio ZIP",
+ )
+ parser.add_argument(
+ "--ravdess-dir",
+ type=Path,
+ default=Path("data/ravdess"),
+ help="Path to RAVDESS directory with manifest.csv",
+ )
+ parser.add_argument(
+ "--output-dir",
+ type=Path,
+ default=Path("data/lora_7class"),
+ help="Output directory for processed dataset",
+ )
+ parser.add_argument("--cap-263", type=int, default=2100, help="Cap per class for 263")
+ parser.add_argument("--cap-71631", type=int, default=3500, help="Cap per class for 71631")
+ parser.add_argument("--max-per-speaker-71631", type=int, default=20, help="Max utterances per speaker for 71631")
+ parser.add_argument("--val-ratio", type=float, default=0.1, help="Val ratio for splits")
+ parser.add_argument("--seed", type=int, default=42, help="Random seed")
+ parser.add_argument(
+ "--skip-263", action="store_true", help="Skip AI Hub 263 extraction"
+ )
+ parser.add_argument(
+ "--skip-71631", action="store_true", help="Skip AI Hub 71631 extraction"
+ )
+ parser.add_argument(
+ "--skip-ravdess", action="store_true", help="Skip RAVDESS extraction"
+ )
+ args = parser.parse_args()
+
+ random.seed(args.seed)
+
+ output_dir = args.output_dir
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ all_train: list[dict] = []
+ all_val: list[dict] = []
+
+ # ---- 263 Anchor ----
+ if not args.skip_263:
+ log.info("=" * 60)
+ log.info("Extracting AI Hub 263 (anchor)")
+ samples_263 = extract_anchor_263(
+ args.anchor_dir,
+ output_dir / "263",
+ cap_per_class=args.cap_263,
+ )
+ train_263, val_263 = speaker_isolated_split(samples_263, args.val_ratio)
+ log.info("263: train=%d, val=%d", len(train_263), len(val_263))
+ all_train.extend(train_263)
+ all_val.extend(val_263)
+
+ # ---- 71631 Booster ----
+ if not args.skip_71631:
+ log.info("=" * 60)
+ log.info("Extracting AI Hub 71631 (booster)")
+ samples_71631 = extract_booster_71631(
+ args.booster_label_zip,
+ args.booster_audio_zip,
+ output_dir / "71631",
+ cap=args.cap_71631,
+ max_per_speaker=args.max_per_speaker_71631,
+ )
+ train_71631, val_71631 = conversation_isolated_split(
+ samples_71631, args.val_ratio
+ )
+ log.info("71631: train=%d, val=%d", len(train_71631), len(val_71631))
+ all_train.extend(train_71631)
+ all_val.extend(val_71631)
+
+ # ---- RAVDESS ----
+ if not args.skip_ravdess:
+ log.info("=" * 60)
+ log.info("Extracting RAVDESS")
+ samples_ravdess = extract_ravdess(
+ args.ravdess_dir,
+ output_dir / "ravdess",
+ )
+ val_actors = [21, 22, 23, 24]
+ train_ravdess, val_ravdess = actor_split_ravdess(
+ samples_ravdess, val_actors
+ )
+ log.info("RAVDESS: train=%d, val=%d", len(train_ravdess), len(val_ravdess))
+ all_train.extend(train_ravdess)
+ all_val.extend(val_ravdess)
+
+ # ---- Save ----
+ log.info("=" * 60)
+ log.info("Total: train=%d, val=%d", len(all_train), len(all_val))
+
+ save_manifest(all_train, output_dir / "train_manifest.json")
+ save_manifest(all_val, output_dir / "val_manifest.json")
+ save_stats(all_train, all_val, output_dir / "stats.json")
+
+ log.info("Done! Output: %s", output_dir)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/prepare_meld_fusion_data.py b/scripts/prepare_meld_fusion_data.py
new file mode 100644
index 0000000000000000000000000000000000000000..3ecf643e79f4668bcf246e677f2f7dc28ee68c6e
--- /dev/null
+++ b/scripts/prepare_meld_fusion_data.py
@@ -0,0 +1,153 @@
+#!/usr/bin/env python3
+"""Prepare MELD test split for English fusion grid search.
+
+Extracts mp4 โ wav (16kHz mono) and builds a manifest with text + emotion labels.
+
+Usage:
+ python scripts/prepare_meld_fusion_data.py
+"""
+from __future__ import annotations
+
+import csv
+import io
+import json
+import logging
+import subprocess
+import tempfile
+import zipfile
+from collections import Counter
+from pathlib import Path
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+PROJECT_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear", "disgust"]
+
+# MELD emotions map 1:1 to project labels
+MELD_LABEL_MAP = {
+ "neutral": "neutral",
+ "joy": "joy",
+ "sadness": "sadness",
+ "anger": "anger",
+ "surprise": "surprise",
+ "fear": "fear",
+ "disgust": "disgust",
+}
+
+
+def main():
+ zip_path = Path("data/english_test.zip")
+ output_dir = Path("data/meld_fusion")
+ audio_dir = output_dir / "audio"
+ audio_dir.mkdir(parents=True, exist_ok=True)
+
+ zf = zipfile.ZipFile(zip_path)
+
+ # Step 1: Parse test CSV
+ logger.info("Parsing MELD test CSV...")
+ with zf.open("MELD.Raw/MELD.Raw/test_sent_emo.csv") as f:
+ reader = csv.DictReader(io.TextIOWrapper(f, encoding="utf-8"))
+ rows = list(reader)
+ logger.info("MELD test: %d utterances", len(rows))
+
+ # Build lookup: (dia_id, utt_id) โ row
+ csv_lookup = {}
+ for r in rows:
+ key = (int(r["Dialogue_ID"]), int(r["Utterance_ID"]))
+ csv_lookup[key] = r
+
+ # Step 2: Find mp4 files in zip
+ test_mp4s = {}
+ for name in zf.namelist():
+ if "output_repeated_splits_test" in name and name.endswith(".mp4"):
+ fname = Path(name).name
+ if fname.startswith("._"):
+ continue # skip macOS metadata
+ # Parse dia{D}_utt{U}.mp4
+ try:
+ parts = fname.replace(".mp4", "").split("_")
+ dia_id = int(parts[0].replace("dia", ""))
+ utt_id = int(parts[1].replace("utt", ""))
+ test_mp4s[(dia_id, utt_id)] = name
+ except (ValueError, IndexError):
+ continue
+
+ logger.info("Found %d test mp4 files (excluding macOS metadata)", len(test_mp4s))
+
+ # Step 3: Match CSV โ mp4, extract wav
+ manifest = []
+ skipped = 0
+
+ matched_keys = set(csv_lookup.keys()) & set(test_mp4s.keys())
+ logger.info("Matched CSVโmp4: %d", len(matched_keys))
+
+ for i, key in enumerate(sorted(matched_keys)):
+ row = csv_lookup[key]
+ mp4_name = test_mp4s[key]
+ dia_id, utt_id = key
+
+ label = MELD_LABEL_MAP.get(row["Emotion"])
+ if label is None:
+ skipped += 1
+ continue
+
+ text = row["Utterance"].strip()
+ if not text:
+ skipped += 1
+ continue
+
+ wav_path = audio_dir / f"dia{dia_id}_utt{utt_id}.wav"
+
+ if not wav_path.exists():
+ # Extract mp4 from zip โ convert to 16kHz mono wav
+ try:
+ mp4_bytes = zf.read(mp4_name)
+ with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
+ tmp.write(mp4_bytes)
+ tmp_path = tmp.name
+
+ result = subprocess.run(
+ ["ffmpeg", "-y", "-i", tmp_path,
+ "-ar", "16000", "-ac", "1", "-f", "wav",
+ str(wav_path)],
+ capture_output=True, timeout=30,
+ )
+ Path(tmp_path).unlink(missing_ok=True)
+
+ if result.returncode != 0:
+ skipped += 1
+ continue
+ except Exception as e:
+ logger.warning("Failed dia%d_utt%d: %s", dia_id, utt_id, e)
+ skipped += 1
+ continue
+
+ manifest.append({
+ "path": str(wav_path),
+ "text": text,
+ "label": label,
+ "source": "meld_test",
+ "dialogue_id": dia_id,
+ "utterance_id": utt_id,
+ })
+
+ if (i + 1) % 200 == 0:
+ logger.info("Processed %d / %d", i + 1, len(matched_keys))
+
+ zf.close()
+
+ # Step 4: Save manifest
+ manifest_path = output_dir / "manifest.json"
+ with open(manifest_path, "w", encoding="utf-8") as f:
+ json.dump(manifest, f, indent=2, ensure_ascii=False)
+
+ logger.info("Saved %d samples to %s (skipped %d)", len(manifest), manifest_path, skipped)
+
+ # Stats
+ emotions = Counter(s["label"] for s in manifest)
+ for e, c in sorted(emotions.items(), key=lambda x: -x[1]):
+ print(f" {e}: {c}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/prepare_ravdess.py b/scripts/prepare_ravdess.py
new file mode 100644
index 0000000000000000000000000000000000000000..7c255a7fab1c332adb622c5afcb2f72de8a1a4cf
--- /dev/null
+++ b/scripts/prepare_ravdess.py
@@ -0,0 +1,220 @@
+#!/usr/bin/env python3
+"""RAVDESS ์์ด ๊ฐ์ ์์ฑ ๋ฐ์ดํฐ์
์ค๋น ์คํฌ๋ฆฝํธ.
+
+data/archive.zip์ ์์ถ ํด์ ํ๊ณ , ์ ํ ํ์ง ์ ์ฒ๋ฆฌ๋ฅผ ์ ์ฉํ์ฌ
+emotion2vec ์์ด ํ๊ฐ์ฉ manifest.csv๋ฅผ ์์ฑํ๋ค.
+
+Usage:
+ python scripts/prepare_ravdess.py
+ python scripts/prepare_ravdess.py --skip-phone # ์ ํ ์ ์ฒ๋ฆฌ ์๋ต
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import logging
+import sys
+import zipfile
+from pathlib import Path
+
+import librosa
+import numpy as np
+import soundfile as sf
+
+PROJECT_ROOT = Path(__file__).parent.parent
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from src.common.phone_simulator import CompandingType, PhoneSimulator
+
+logging.basicConfig(
+ level=logging.INFO,
+ format="%(asctime)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger("prepare_ravdess")
+
+# RAVDESS emotion code โ project 7-class taxonomy
+RAVDESS_EMOTION_MAP = {
+ 1: "neutral", # neutral
+ 2: "neutral", # calm โ neutral (ํ๋ก์ ํธ taxonomy์ calm ์์)
+ 3: "joy", # happy โ joy
+ 4: "sadness", # sad โ sadness
+ 5: "anger", # angry โ anger
+ 6: "fear", # fearful โ fear
+ 7: "disgust", # disgust
+ 8: "surprise", # surprised โ surprise
+}
+
+RAVDESS_EMOTION_NAME = {
+ 1: "neutral", 2: "calm", 3: "happy", 4: "sad",
+ 5: "angry", 6: "fearful", 7: "disgust", 8: "surprised",
+}
+
+ARCHIVE_PATH = PROJECT_ROOT / "data" / "archive.zip"
+OUTPUT_DIR = PROJECT_ROOT / "data" / "ravdess"
+
+
+def parse_ravdess_filename(filename: str) -> dict | None:
+ """RAVDESS ํ์ผ๋ช
์์ ๋ฉํ๋ฐ์ดํฐ ์ถ์ถ.
+
+ Format: Modality-VocalChannel-Emotion-Intensity-Statement-Repetition-Actor.wav
+ Example: 03-01-05-02-01-01-12.wav
+ """
+ stem = Path(filename).stem
+ parts = stem.split("-")
+ if len(parts) != 7:
+ return None
+
+ emotion_code = int(parts[2])
+ return {
+ "modality": int(parts[0]),
+ "vocal_channel": int(parts[1]),
+ "emotion_code": emotion_code,
+ "emotion_raw": RAVDESS_EMOTION_NAME.get(emotion_code, "unknown"),
+ "emotion": RAVDESS_EMOTION_MAP.get(emotion_code, "neutral"),
+ "intensity": int(parts[3]), # 1=normal, 2=strong
+ "statement": int(parts[4]), # 1="Kids...", 2="Dogs..."
+ "repetition": int(parts[5]),
+ "actor_id": int(parts[6]),
+ }
+
+
+def extract_archive(archive_path: Path, output_dir: Path) -> list[Path]:
+ """archive.zip ์์ถ ํด์ โ clean/ ๋๋ ํ ๋ฆฌ."""
+ clean_dir = output_dir / "clean"
+
+ if clean_dir.exists() and any(clean_dir.rglob("*.wav")):
+ wavs = sorted(clean_dir.rglob("*.wav"))
+ logger.info(f"์ด๋ฏธ ์์ถ ํด์ ๋จ: {len(wavs)}๊ฐ WAV in {clean_dir}")
+ return wavs
+
+ clean_dir.mkdir(parents=True, exist_ok=True)
+ logger.info(f"์์ถ ํด์ ์ค: {archive_path} โ {clean_dir}")
+
+ with zipfile.ZipFile(archive_path, "r") as zf:
+ wav_members = [m for m in zf.namelist() if m.endswith(".wav")]
+ for i, member in enumerate(wav_members, 1):
+ # Actor_NN/filename.wav โ clean/Actor_NN/filename.wav
+ target = clean_dir / member
+ target.parent.mkdir(parents=True, exist_ok=True)
+ with zf.open(member) as src, open(target, "wb") as dst:
+ dst.write(src.read())
+ if i % 500 == 0:
+ logger.info(f" [{i}/{len(wav_members)}] ์์ถ ํด์ ์ค...")
+
+ wavs = sorted(clean_dir.rglob("*.wav"))
+ logger.info(f"์์ถ ํด์ ์๋ฃ: {len(wavs)}๊ฐ WAV")
+ return wavs
+
+
+def apply_phone_simulation(clean_wavs: list[Path], output_dir: Path) -> dict[str, Path]:
+ """clean WAV โ phone ํ์ง ๋ณํ. {clean_path_str: phone_path} ๋ฐํ."""
+ phone_dir = output_dir / "phone"
+ simulator = PhoneSimulator(companding=CompandingType.ULAW) # ์์ด = ๋ถ๋ฏธ ฮผ-law
+
+ mapping = {}
+ total = len(clean_wavs)
+
+ for i, wav_path in enumerate(clean_wavs, 1):
+ # clean/Actor_NN/file.wav โ phone/Actor_NN/file.wav
+ relative = wav_path.relative_to(output_dir / "clean")
+ phone_path = phone_dir / relative
+ phone_path.parent.mkdir(parents=True, exist_ok=True)
+
+ if phone_path.exists():
+ mapping[str(wav_path)] = phone_path
+ continue
+
+ try:
+ audio, sr = librosa.load(str(wav_path), sr=None, mono=True)
+ processed, new_sr = simulator.process(audio, sr)
+ sf.write(str(phone_path), processed, new_sr, subtype="PCM_16")
+ mapping[str(wav_path)] = phone_path
+ except Exception as e:
+ logger.warning(f"์ ํ ๋ณํ ์คํจ [{wav_path.name}]: {e}")
+
+ if i % 500 == 0:
+ logger.info(f" [{i}/{total}] ์ ํ ํ์ง ๋ณํ ์ค...")
+
+ logger.info(f"์ ํ ํ์ง ๋ณํ ์๋ฃ: {len(mapping)}/{total}")
+ return mapping
+
+
+def build_manifest(
+ clean_wavs: list[Path],
+ phone_mapping: dict[str, Path] | None,
+ output_dir: Path,
+) -> Path:
+ """manifest.csv ์์ฑ."""
+ manifest_path = output_dir / "manifest.csv"
+ rows = []
+
+ for wav_path in clean_wavs:
+ meta = parse_ravdess_filename(wav_path.name)
+ if meta is None:
+ logger.warning(f"ํ์ผ๋ช
ํ์ฑ ์คํจ: {wav_path.name}")
+ continue
+
+ phone_path = ""
+ if phone_mapping and str(wav_path) in phone_mapping:
+ phone_path = str(phone_mapping[str(wav_path)])
+
+ rows.append({
+ "clean_path": str(wav_path),
+ "phone_path": phone_path,
+ "emotion": meta["emotion"],
+ "emotion_raw": meta["emotion_raw"],
+ "actor_id": meta["actor_id"],
+ "intensity": meta["intensity"],
+ "statement": meta["statement"],
+ "repetition": meta["repetition"],
+ })
+
+ # ๊ฐ์ ๋ณ ํต๊ณ ์ถ๋ ฅ
+ from collections import Counter
+ emotion_counts = Counter(r["emotion"] for r in rows)
+ logger.info("๊ฐ์ ๋ถํฌ:")
+ for emotion, count in sorted(emotion_counts.items()):
+ logger.info(f" {emotion}: {count}")
+
+ with open(manifest_path, "w", newline="") as f:
+ writer = csv.DictWriter(f, fieldnames=[
+ "clean_path", "phone_path", "emotion", "emotion_raw",
+ "actor_id", "intensity", "statement", "repetition",
+ ])
+ writer.writeheader()
+ writer.writerows(rows)
+
+ logger.info(f"manifest ์ ์ฅ: {manifest_path} ({len(rows)}ํ)")
+ return manifest_path
+
+
+def main():
+ parser = argparse.ArgumentParser(description="RAVDESS ์์ด ๊ฐ์ ๋ฐ์ดํฐ ์ค๋น")
+ parser.add_argument("--skip-phone", action="store_true", help="์ ํ ํ์ง ์ ์ฒ๋ฆฌ ์๋ต")
+ args = parser.parse_args()
+
+ if not ARCHIVE_PATH.exists():
+ logger.error(f"archive.zip์ ์ฐพ์ ์ ์์ต๋๋ค: {ARCHIVE_PATH}")
+ sys.exit(1)
+
+ OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+ # 1. ์์ถ ํด์
+ clean_wavs = extract_archive(ARCHIVE_PATH, OUTPUT_DIR)
+
+ # 2. ์ ํ ํ์ง ์ ์ฒ๋ฆฌ
+ phone_mapping = None
+ if not args.skip_phone:
+ phone_mapping = apply_phone_simulation(clean_wavs, OUTPUT_DIR)
+ else:
+ logger.info("์ ํ ํ์ง ์ ์ฒ๋ฆฌ ์๋ต (--skip-phone)")
+
+ # 3. manifest.csv ์์ฑ
+ build_manifest(clean_wavs, phone_mapping, OUTPUT_DIR)
+
+ logger.info("์๋ฃ!")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/preprocess_phone_audio.py b/scripts/preprocess_phone_audio.py
new file mode 100644
index 0000000000000000000000000000000000000000..3be28010cf31582e4829d506d203b760413ea914
--- /dev/null
+++ b/scripts/preprocess_phone_audio.py
@@ -0,0 +1,175 @@
+#!/usr/bin/env python3
+"""AI Hub ๋ฑ ์คํ๋์ค ๋
น์ ๋ฐ์ดํฐ๋ฅผ ์ ํ ํตํ ํ์ง๋ก ์ ์ฒ๋ฆฌํ๋ ์คํฌ๋ฆฝํธ.
+
+๊นจ๋ํ ์ค๋์ค์ PSTN ์๋ฎฌ๋ ์ด์
(๋ฐด๋ํจ์ค + ๋ค์ด์ํ๋ง + G.711 companding)์ ์ ์ฉํ์ฌ
+์ค์ ํตํ ๋
น์๊ณผ ์ ์ฌํ ํ์ต ๋ฐ์ดํฐ๋ฅผ ์์ฑํ๋ค.
+
+Usage:
+ # ๋จ์ผ ํ์ผ
+ python scripts/preprocess_phone_audio.py data/aihub_raw/sample.wav
+
+ # ๋๋ ํ ๋ฆฌ ์ผ๊ด ์ฒ๋ฆฌ
+ python scripts/preprocess_phone_audio.py data/aihub_raw/ -o data/aihub_phone/
+
+ # companding ๋ฐฉ์ ์ง์ (๊ธฐ๋ณธ: random)
+ python scripts/preprocess_phone_audio.py data/aihub_raw/ --companding alaw
+
+ # ์๋ณธ๋ ํจ๊ป ๋ณต์ฌ (์๋ณธ+์ ํ ํผํฉ ํ์ต์ฉ)
+ python scripts/preprocess_phone_audio.py data/aihub_raw/ -o data/training/ --keep-original
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import shutil
+import sys
+from pathlib import Path
+
+import librosa
+import soundfile as sf
+
+# ํ๋ก์ ํธ ๋ฃจํธ๋ฅผ sys.path์ ์ถ๊ฐ
+PROJECT_ROOT = Path(__file__).parent.parent
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from src.common.phone_simulator import CompandingType, PhoneSimulator
+
+logging.basicConfig(
+ level=logging.INFO,
+ format="%(asctime)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger("preprocess_phone_audio")
+
+SUPPORTED_EXTENSIONS = {".wav", ".mp3", ".m4a", ".ogg", ".flac"}
+
+
+def find_audio_files(input_path: Path) -> list[Path]:
+ """์
๋ ฅ ๊ฒฝ๋ก์์ ์ค๋์ค ํ์ผ ๋ชฉ๋ก ๋ฐํ."""
+ if input_path.is_file():
+ if input_path.suffix.lower() in SUPPORTED_EXTENSIONS:
+ return [input_path]
+ logger.warning(f"์ง์ํ์ง ์๋ ํ์ผ ํ์: {input_path.suffix}")
+ return []
+
+ files = []
+ for ext in SUPPORTED_EXTENSIONS:
+ files.extend(input_path.rglob(f"*{ext}"))
+ return sorted(files)
+
+
+def process_file(
+ input_file: Path,
+ output_dir: Path,
+ simulator: PhoneSimulator,
+ input_root: Path,
+ keep_original: bool = False,
+) -> bool:
+ """๋จ์ผ ํ์ผ์ ์ ํ ํ์ง๋ก ๋ณํ."""
+ try:
+ # ์๋ณธ ๋๋ ํ ๋ฆฌ ๊ตฌ์กฐ ์ ์ง
+ relative = input_file.relative_to(input_root)
+ output_file = output_dir / relative.with_suffix(".wav")
+ output_file.parent.mkdir(parents=True, exist_ok=True)
+
+ # ์ค๋์ค ๋ก๋ (mono, ์๋ณธ SR ์ ์ง)
+ audio, sr = librosa.load(str(input_file), sr=None, mono=True)
+
+ # ์ ํ ํ์ง ์๋ฎฌ๋ ์ด์
์ ์ฉ
+ processed, new_sr = simulator.process(audio, sr)
+
+ # phone_ ์ ๋์ฌ๋ก ์ ์ฅ
+ phone_output = output_file.with_name(f"phone_{output_file.name}")
+ sf.write(str(phone_output), processed, new_sr, subtype="PCM_16")
+
+ # ์๋ณธ๋ ๋ณต์ฌ (ํผํฉ ํ์ต์ฉ)
+ if keep_original:
+ orig_output = output_file.with_name(f"orig_{output_file.name}")
+ shutil.copy2(str(input_file), str(orig_output))
+
+ return True
+
+ except Exception as e:
+ logger.error(f"์ฒ๋ฆฌ ์คํจ [{input_file.name}]: {e}")
+ return False
+
+
+def main():
+ parser = argparse.ArgumentParser(
+ description="์คํ๋์ค ๋
น์ โ ์ ํ ํตํ ํ์ง ์ ์ฒ๋ฆฌ",
+ )
+ parser.add_argument(
+ "input",
+ type=Path,
+ help="์
๋ ฅ ์ค๋์ค ํ์ผ ๋๋ ๋๋ ํ ๋ฆฌ",
+ )
+ parser.add_argument(
+ "-o", "--output",
+ type=Path,
+ default=None,
+ help="์ถ๋ ฅ ๋๋ ํ ๋ฆฌ (๊ธฐ๋ณธ: {input}_phone/)",
+ )
+ parser.add_argument(
+ "--companding",
+ type=str,
+ choices=["alaw", "ulaw", "random"],
+ default="random",
+ help="G.711 companding ๋ฐฉ์ (๊ธฐ๋ณธ: random โ ํ์ผ๋ง๋ค ๋๋ค ์ ํ)",
+ )
+ parser.add_argument(
+ "--keep-original",
+ action="store_true",
+ help="์๋ณธ ํ์ผ๋ ์ถ๋ ฅ ๋๋ ํ ๋ฆฌ์ ๋ณต์ฌ (์๋ณธ+์ ํ ํผํฉ ํ์ต์ฉ)",
+ )
+ args = parser.parse_args()
+
+ # ์
๋ ฅ ๊ฒฝ๋ก ํ์ธ
+ input_path = args.input.resolve()
+ if not input_path.exists():
+ logger.error(f"์
๋ ฅ ๊ฒฝ๋ก๊ฐ ์กด์ฌํ์ง ์์ต๋๋ค: {input_path}")
+ sys.exit(1)
+
+ # ์ถ๋ ฅ ๋๋ ํ ๋ฆฌ ๊ฒฐ์
+ if args.output:
+ output_dir = args.output.resolve()
+ else:
+ if input_path.is_file():
+ output_dir = input_path.parent / f"{input_path.stem}_phone"
+ else:
+ output_dir = input_path.parent / f"{input_path.name}_phone"
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ # ์
๋ ฅ ๋ฃจํธ (์๋ ๊ฒฝ๋ก ๊ณ์ฐ์ฉ)
+ input_root = input_path if input_path.is_dir() else input_path.parent
+
+ # ์ค๋์ค ํ์ผ ํ์
+ audio_files = find_audio_files(input_path)
+ if not audio_files:
+ logger.error("์ฒ๋ฆฌํ ์ค๋์ค ํ์ผ์ด ์์ต๋๋ค.")
+ sys.exit(1)
+
+ logger.info(f"์ค๋์ค ํ์ผ {len(audio_files)}๊ฐ ๋ฐ๊ฒฌ")
+ logger.info(f"์ถ๋ ฅ ๋๋ ํ ๋ฆฌ: {output_dir}")
+ logger.info(f"Companding: {args.companding}")
+ if args.keep_original:
+ logger.info("์๋ณธ ํ์ผ๋ ํจ๊ป ๋ณต์ฌํฉ๋๋ค")
+
+ # ์๋ฎฌ๋ ์ดํฐ ์์ฑ
+ companding = CompandingType(args.companding)
+ simulator = PhoneSimulator(companding=companding)
+
+ # ์ผ๊ด ์ฒ๋ฆฌ
+ success = 0
+ fail = 0
+ for i, audio_file in enumerate(audio_files, 1):
+ logger.info(f"[{i}/{len(audio_files)}] {audio_file.name}")
+ if process_file(audio_file, output_dir, simulator, input_root, args.keep_original):
+ success += 1
+ else:
+ fail += 1
+
+ logger.info(f"์๋ฃ: ์ฑ๊ณต {success}, ์คํจ {fail}, ์ ์ฒด {len(audio_files)}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/quantize_model.py b/scripts/quantize_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/scripts/run_pipeline.py b/scripts/run_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..033015239bae1d0ebfa18d50534759dd0c27c5cc
--- /dev/null
+++ b/scripts/run_pipeline.py
@@ -0,0 +1,121 @@
+#!/usr/bin/env python3
+"""Stage 1 โ Stage 2 E2E ํ์ดํ๋ผ์ธ ์คํ ์คํฌ๋ฆฝํธ.
+
+Usage:
+ python scripts/run_pipeline.py # ๊ธฐ๋ณธ ์ํ ์ฌ์ฉ
+ python scripts/run_pipeline.py data/samples/my_call.wav # ํน์ ํ์ผ ์ง์
+ python scripts/run_pipeline.py --stage2-only # Stage 2๋ง ์คํ (stage1_output.json ํ์)
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import sys
+from pathlib import Path
+
+# ํ๋ก์ ํธ ๋ฃจํธ๋ฅผ sys.path์ ์ถ๊ฐ
+PROJECT_ROOT = Path(__file__).parent.parent
+sys.path.insert(0, str(PROJECT_ROOT))
+
+logging.basicConfig(
+ level=logging.INFO,
+ format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger("run_pipeline")
+
+
+def run_stage1(audio_path: str) -> dict:
+ """Stage 1 ์คํ: ํ์๋ถ๋ฆฌ + ASR."""
+ from src.stage1.process import process as stage1_process
+
+ logger.info("=" * 60)
+ logger.info("Stage 1 ์์: %s", audio_path)
+ logger.info("=" * 60)
+
+ result = stage1_process(audio_path)
+
+ logger.info(
+ "Stage 1 ์๋ฃ: %d segments, %.1fs ์ฒ๋ฆฌ์๊ฐ",
+ len(result.segments),
+ result.processing_info.processing_time_sec,
+ )
+ return result
+
+
+def run_stage2(stage1_output) -> dict:
+ """Stage 2 ์คํ: ๊ฐ์ ๋ถ์."""
+ from src.stage2.process import process as stage2_process
+
+ logger.info("=" * 60)
+ logger.info("Stage 2 ์์: %s (%d segments)", stage1_output.call_id, len(stage1_output.segments))
+ logger.info("=" * 60)
+
+ result = stage2_process(stage1_output)
+
+ logger.info(
+ "Stage 2 ์๋ฃ: %d emotions, speakers=%s",
+ len(result.emotions),
+ list(result.speaker_summaries.keys()),
+ )
+
+ # ๊ฒฐ๊ณผ ์์ฝ ์ถ๋ ฅ
+ for speaker_id, summary in result.speaker_summaries.items():
+ logger.info(
+ " %s: dominant=%s (%.1f%%), avg_confidence=%.2f",
+ speaker_id,
+ summary.dominant_emotion,
+ summary.emotion_distribution.get(summary.dominant_emotion, 0) * 100,
+ summary.avg_confidence,
+ )
+
+ return result
+
+
+def main():
+ parser = argparse.ArgumentParser(description="Stage 1 โ Stage 2 ํ์ดํ๋ผ์ธ ์คํ")
+ parser.add_argument(
+ "audio_path",
+ nargs="?",
+ default="data/samples/sample_data.wav",
+ help="์
๋ ฅ ์ค๋์ค ํ์ผ ๊ฒฝ๋ก (๊ธฐ๋ณธ: data/samples/sample_data.wav)",
+ )
+ parser.add_argument(
+ "--stage2-only",
+ action="store_true",
+ help="Stage 2๋ง ์คํ (data/stage1_output.json ํ์)",
+ )
+ args = parser.parse_args()
+
+ if args.stage2_only:
+ # Stage 2๋ง ์คํ
+ from src.common.schemas import Stage1Output
+
+ stage1_path = PROJECT_ROOT / "data" / "stage1_output.json"
+ if not stage1_path.exists():
+ logger.error("data/stage1_output.json ์์. Stage 1์ ๋จผ์ ์คํํ์ธ์.")
+ sys.exit(1)
+
+ stage1_output = Stage1Output.model_validate_json(stage1_path.read_text())
+ run_stage2(stage1_output)
+ else:
+ # Stage 1 โ Stage 2 ์ ์ฒด ์คํ
+ audio_path = str(PROJECT_ROOT / args.audio_path) if not Path(args.audio_path).is_absolute() else args.audio_path
+
+ if not Path(audio_path).exists():
+ logger.error("์ค๋์ค ํ์ผ ์์: %s", audio_path)
+ sys.exit(1)
+
+ stage1_output = run_stage1(audio_path)
+ run_stage2(stage1_output)
+
+ logger.info("=" * 60)
+ logger.info("ํ์ดํ๋ผ์ธ ์๋ฃ!")
+ logger.info(" Stage 1 ์ถ๋ ฅ: data/stage1_output.json")
+ logger.info(" Stage 2 ์ถ๋ ฅ: data/stage2_output.json")
+ logger.info("=" * 60)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/test_20hours_e2e_server.py b/scripts/test_20hours_e2e_server.py
new file mode 100644
index 0000000000000000000000000000000000000000..234894b4176524149a0b36cac82fb097768a4c1a
--- /dev/null
+++ b/scripts/test_20hours_e2e_server.py
@@ -0,0 +1,218 @@
+#!/usr/bin/env python3
+"""20Hours Korean demo set โ E2E server test.
+
+Uploads 7 curated Korean demo WAVs to the deployed HF Spaces server,
+runs the full pipeline (Stage 1 โ 2 โ 3), and compares results
+against the intended demo emotion labels (data/20hours_test/ground_truth.json).
+
+Usage:
+ python scripts/test_20hours_e2e_server.py
+ python scripts/test_20hours_e2e_server.py --server http://localhost:8000
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+import requests
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+TEST_DIR = PROJECT_ROOT / "data" / "20hours_test"
+GT_PATH = TEST_DIR / "ground_truth.json"
+
+DEFAULT_SERVER = "https://bbbakery-ustwo-api.hf.space"
+POLL_INTERVAL = 5
+MAX_WAIT = 300
+
+
+def health_check(base: str) -> bool:
+ try:
+ r = requests.get(f"{base}/api/health", timeout=10)
+ data = r.json()
+ if data.get("status") == "ok":
+ print(f" Server OK ({data.get('timestamp', '?')})")
+ return True
+ except Exception as e:
+ print(f" Health check failed: {e}")
+ return False
+
+
+def upload(base: str, wav_path: Path) -> str | None:
+ with open(wav_path, "rb") as f:
+ r = requests.post(
+ f"{base}/api/upload",
+ files={"file": (wav_path.name, f, "audio/wav")},
+ timeout=60,
+ )
+ if r.status_code != 200:
+ print(f" Upload failed ({r.status_code}): {r.text[:200]}")
+ return None
+ return r.json().get("call_id")
+
+
+def analyze_and_poll(base: str, call_id: str) -> dict | None:
+ r = requests.post(f"{base}/api/analyze", params={"call_id": call_id}, timeout=30)
+ if r.status_code not in (200, 202):
+ print(f" Analyze start failed ({r.status_code}): {r.text[:200]}")
+ return None
+
+ data = r.json()
+ if data.get("status") == "done":
+ return data.get("result")
+
+ elapsed = 0
+ while elapsed < MAX_WAIT:
+ time.sleep(POLL_INTERVAL)
+ elapsed += POLL_INTERVAL
+ r = requests.get(f"{base}/api/analyze/{call_id}/status", timeout=15)
+ data = r.json()
+ status = data.get("status")
+ if status == "done":
+ return data.get("result")
+ if status == "error":
+ print(f" Pipeline error: {data.get('error', '?')}")
+ return None
+ mins, secs = divmod(elapsed, 60)
+ print(f" {status}... ({int(mins)}m{int(secs)}s)", end="\r")
+
+ print(f" Timeout after {MAX_WAIT}s")
+ return None
+
+
+def extract_emotions(result: dict) -> dict:
+ info: dict = {}
+ reactions = result.get("character_reactions", [])
+ for i, rx in enumerate(reactions):
+ info[f"speaker_{i}_state"] = rx.get("solo_state", "?")
+ garden = result.get("garden_update", {})
+ info["garden_mood"] = garden.get("mood", "?")
+ info["garden_delta"] = garden.get("growth_delta", 0)
+ recap = result.get("recap_card", {}) or {}
+ info["recap_headline"] = recap.get("headline") or recap.get("title", "?")
+
+ stage2 = result.get("stage2_output", {})
+ if stage2:
+ for spk, summary in stage2.get("speaker_summaries", {}).items():
+ info[f"{spk}_dominant"] = summary.get("dominant_emotion", "?")
+ info[f"{spk}_distribution"] = summary.get("emotion_distribution", {})
+
+ # Segment-level language breakdown
+ emotions = result.get("emotions") or stage2.get("emotions", [])
+ segs_by_lang: dict[str, int] = {}
+ for e in emotions:
+ lang = e.get("language") or "?"
+ segs_by_lang[lang] = segs_by_lang.get(lang, 0) + 1
+ info["segments"] = len(emotions)
+ info["segments_by_lang"] = segs_by_lang
+ return info
+
+
+def main():
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--server", default=DEFAULT_SERVER)
+ args = parser.parse_args()
+ base = args.server.rstrip("/")
+
+ print("=" * 70)
+ print(" 20Hours Korean Demo โ E2E Server Test")
+ print(f" Server: {base}")
+ print("=" * 70)
+
+ print("\n[1] Health check")
+ if not health_check(base):
+ sys.exit(1)
+
+ print("\n[2] Loading intended emotion labels")
+ gt = json.loads(GT_PATH.read_text())
+ print(f" {len(gt)} demo clips loaded")
+
+ print("\n[3] Running E2E tests\n")
+ results = {}
+ hit = 0
+ total = 0
+
+ for tag in sorted(gt.keys()):
+ wav_path = TEST_DIR / f"{tag}.wav"
+ if not wav_path.exists():
+ print(f" {tag}: WAV not found, skipping")
+ continue
+ gt_entry = gt[tag]
+ print(f" {tag} โ {gt_entry['description'][:55]}")
+ print(f" Intended: {gt_entry['primary_emotion']} | Duration: {gt_entry['duration_sec']}s | Utts: {gt_entry['total_utterances']}")
+
+ call_id = upload(base, wav_path)
+ if not call_id:
+ results[tag] = {"status": "upload_failed"}
+ continue
+ print(f" Upload OK โ {call_id}")
+
+ print(f" Analyzing...", end="")
+ start_time = time.time()
+ result = analyze_and_poll(base, call_id)
+ elapsed = time.time() - start_time
+
+ if not result:
+ results[tag] = {"status": "analyze_failed", "call_id": call_id}
+ print()
+ continue
+
+ print(f"\r Done in {elapsed:.1f}s ")
+
+ emotions = extract_emotions(result)
+ total += 1
+ intended = gt_entry["primary_emotion"]
+ speaker_states = {k: v for k, v in emotions.items() if k.endswith("_state")}
+ if intended in speaker_states.values():
+ hit += 1
+ match = "HIT"
+ else:
+ match = "miss"
+
+ results[tag] = {
+ "status": "pass",
+ "call_id": call_id,
+ "elapsed_sec": round(elapsed, 1),
+ "intended_emotion": intended,
+ "match": match,
+ "emotions": emotions,
+ "full_result": result,
+ }
+
+ for k, v in emotions.items():
+ if not k.endswith("_distribution"):
+ print(f" {k}: {v}")
+ else:
+ dist = ", ".join(f"{kk}:{vv:.2f}" for kk, vv in sorted(v.items(), key=lambda x: -x[1])[:3])
+ print(f" {k}: {dist}")
+ print(f" โ {match}")
+ print()
+
+ out_path = TEST_DIR / "e2e_results.json"
+ save = {}
+ for tag, r in results.items():
+ save[tag] = {k: v for k, v in r.items() if k != "full_result"}
+ out_path.write_text(json.dumps(save, indent=2, ensure_ascii=False))
+
+ print("=" * 70)
+ print(" SUMMARY")
+ print("=" * 70)
+ print(f"\n {'Tag':<24} {'Intended':<12} {'Match':<6} {'Time':>6}")
+ print(" " + "-" * 55)
+ for tag in sorted(results.keys()):
+ r = results[tag]
+ if r.get("status") != "pass":
+ continue
+ intended = r["intended_emotion"]
+ m = r["match"]
+ t = f"{r['elapsed_sec']:.0f}s"
+ print(f" {tag:<24} {intended:<12} {m:<6} {t:>6}")
+ print(f"\n Intended-emotion match: {hit}/{total}")
+ print(f" Results saved: {out_path}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/test_english_e2e.py b/scripts/test_english_e2e.py
new file mode 100644
index 0000000000000000000000000000000000000000..dd45c87c27a392364494d14f05adbeabf0b85b64
--- /dev/null
+++ b/scripts/test_english_e2e.py
@@ -0,0 +1,239 @@
+#!/usr/bin/env python3
+"""์์ด E2E ํ์ดํ๋ผ์ธ ํ
์คํธ.
+
+RAVDESS WAV + ์๋ ์์ฑ ํ
์คํธ๋ก Stage2 ์ ์ฒด ํ์ดํ๋ผ์ธ์ ํ
์คํธํ๋ค.
+๊ฐ์ ๋ณ 2๊ฐ์ฉ 14๊ฐ ์ธ๊ทธ๋จผํธ๋ฅผ ํต๊ณผ์์ผ audio + text + fusion ๋์์ ๊ฒ์ฆ.
+
+Usage:
+ python scripts/test_english_e2e.py
+"""
+
+from __future__ import annotations
+
+import csv
+import json
+import logging
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).parent.parent
+sys.path.insert(0, str(PROJECT_ROOT))
+
+logging.basicConfig(
+ level=logging.INFO,
+ format="%(asctime)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger("test_english_e2e")
+
+MANIFEST_PATH = PROJECT_ROOT / "data" / "ravdess" / "manifest.csv"
+
+# ๊ฐ์ ๋ณ ๋งค์นญ ํ
์คํธ (RAVDESS๋ ํ
์คํธ ์์ผ๋ฏ๋ก ์๋ ์์ฑ)
+EMOTION_TEXTS = {
+ "neutral": [
+ "The meeting is scheduled for three o'clock.",
+ "I'll pick up the groceries on the way home.",
+ ],
+ "joy": [
+ "I'm so happy to hear from you!",
+ "That's wonderful news, I'm thrilled!",
+ ],
+ "sadness": [
+ "I miss you so much, it hurts.",
+ "I feel really down today, nothing is going right.",
+ ],
+ "anger": [
+ "I can't believe you did that, I'm so angry!",
+ "This is completely unacceptable, stop it now!",
+ ],
+ "surprise": [
+ "Oh my god, I can't believe it!",
+ "What?! I never expected that!",
+ ],
+ "fear": [
+ "I'm scared, something doesn't feel right.",
+ "Please help, I'm terrified right now.",
+ ],
+ "disgust": [
+ "That's absolutely disgusting, I feel sick.",
+ "This is revolting, I can't stand it.",
+ ],
+}
+
+
+def pick_samples(manifest_path: Path) -> list[dict]:
+ """๊ฐ์ ๋ณ 2๊ฐ์ฉ WAV ์ ํ."""
+ by_emotion: dict[str, list[dict]] = {}
+ with open(manifest_path) as f:
+ for row in csv.DictReader(f):
+ emotion = row["emotion"]
+ if emotion not in by_emotion:
+ by_emotion[emotion] = []
+ if len(by_emotion[emotion]) < 2:
+ by_emotion[emotion].append(row)
+
+ samples = []
+ for emotion in EMOTION_TEXTS:
+ if emotion in by_emotion:
+ samples.extend(by_emotion[emotion])
+ return samples
+
+
+def run_e2e():
+ """E2E ํ
์คํธ ์คํ."""
+ from src.common.schemas import (
+ Models,
+ ProcessingInfo,
+ Segment,
+ Stage1Output,
+ )
+ from src.stage2.process import process
+
+ if not MANIFEST_PATH.exists():
+ logger.error("manifest.csv๋ฅผ ์ฐพ์ ์ ์์ต๋๋ค. ๋จผ์ prepare_ravdess.py๋ฅผ ์คํํ์ธ์.")
+ return False
+
+ samples = pick_samples(MANIFEST_PATH)
+ logger.info(f"์ ํ๋ ํ
์คํธ ์ํ: {len(samples)}๊ฐ")
+
+ # Stage1Output ์์ฑ
+ segments = []
+ ground_truths = []
+ for i, sample in enumerate(samples):
+ emotion = sample["emotion"]
+ texts = EMOTION_TEXTS[emotion]
+ text = texts[i % len(texts)]
+
+ segments.append(Segment(
+ segment_id=i,
+ speaker_id=f"speaker_{int(sample['actor_id']) % 2}",
+ start=float(i * 3.0),
+ end=float(i * 3.0 + 3.0),
+ text=text,
+ language="en",
+ audio_path=sample["clean_path"],
+ confidence=0.95,
+ ))
+ ground_truths.append(emotion)
+
+ stage1_output = Stage1Output(
+ call_id="ravdess_e2e_test",
+ duration=float(len(segments) * 3.0),
+ speakers=["speaker_0", "speaker_1"],
+ audio_path="data/ravdess/test_call.wav",
+ segments=segments,
+ processing_info=ProcessingInfo(
+ processing_time_sec=1.0,
+ models=Models(
+ diarization="pyannote/speaker-diarization-3.1",
+ asr="large-v3-turbo",
+ language_id="whisper",
+ alignment="whisperx",
+ ),
+ device="cpu",
+ ),
+ )
+
+ # Stage2 config (์์ด ์ ์ฉ)
+ config = {
+ "audio_emotion": {
+ "model": "iic/emotion2vec_plus_base",
+ },
+ "text_emotion": {
+ "korean_model": "searle-j/kote_for_easygoing_people",
+ "english_model": "j-hartmann/emotion-english-distilroberta-base",
+ },
+ "fusion": {
+ "audio_weight": 0.6,
+ "text_weight": 0.4,
+ },
+ "output_path": "data/e2e_test_output.json",
+ }
+
+ logger.info("Stage 2 process() ์คํ ์ค...")
+ output = process(stage1_output, config=config)
+
+ # ๊ฒ์ฆ
+ logger.info("\n" + "=" * 70)
+ logger.info("E2E ํ
์คํธ ๊ฒฐ๊ณผ")
+ logger.info("=" * 70)
+
+ errors = []
+
+ # 1. EmotionResult ๊ฐ์ ํ์ธ
+ if len(output.emotions) != len(segments):
+ errors.append(f"EmotionResult ๊ฐ์ ๋ถ์ผ์น: {len(output.emotions)} != {len(segments)}")
+ else:
+ logger.info(f"[PASS] EmotionResult ๊ฐ์: {len(output.emotions)}")
+
+ # 2. SpeakerSummary ํ์ธ
+ if len(output.speaker_summaries) > 0:
+ logger.info(f"[PASS] SpeakerSummary ์์ฑ: {len(output.speaker_summaries)}๋ช
")
+ else:
+ errors.append("SpeakerSummary๊ฐ ๋น์ด์์")
+
+ # 3. ๊ฐ EmotionResult ์ ํจ์ฑ ํ์ธ
+ valid_emotions = {"neutral", "joy", "sadness", "anger", "surprise", "fear", "disgust"}
+ for em in output.emotions:
+ if em.fused_emotion not in valid_emotions:
+ errors.append(f"์ ํจํ์ง ์์ fused_emotion: {em.fused_emotion}")
+ if not (0.0 <= em.fused_confidence <= 1.0):
+ errors.append(f"fused_confidence ๋ฒ์ ์ด๊ณผ: {em.fused_confidence}")
+
+ if not any("์ ํจํ์ง ์์" in e for e in errors):
+ logger.info("[PASS] ๋ชจ๋ EmotionResult ์ ํจ")
+
+ # 4. JSON ์ง๋ ฌํ ํ
์คํธ
+ try:
+ json_str = output.model_dump_json(indent=2)
+ roundtrip = json.loads(json_str)
+ logger.info(f"[PASS] JSON ์ง๋ ฌํ/์ญ์ง๋ ฌํ ์ฑ๊ณต ({len(json_str)} bytes)")
+ except Exception as e:
+ errors.append(f"JSON ์ง๋ ฌํ ์คํจ: {e}")
+
+ # 5. ๊ฒฐ๊ณผ ํ
์ด๋ธ
+ logger.info(f"\n{'Seg':>3s} | {'Ground Truth':>12s} | {'Audio':>10s} | {'Text':>10s} | {'Fused':>10s} | {'Conf':>5s}")
+ logger.info("-" * 65)
+ audio_correct = 0
+ text_correct = 0
+ fused_correct = 0
+ for i, (em, gt) in enumerate(zip(output.emotions, ground_truths)):
+ a_match = "o" if em.audio_emotion == gt else "x"
+ t_match = "o" if em.text_emotion == gt else "x"
+ f_match = "o" if em.fused_emotion == gt else "x"
+ if em.audio_emotion == gt:
+ audio_correct += 1
+ if em.text_emotion == gt:
+ text_correct += 1
+ if em.fused_emotion == gt:
+ fused_correct += 1
+ logger.info(
+ f"{i:3d} | {gt:>12s} | {em.audio_emotion:>8s} {a_match} | "
+ f"{em.text_emotion:>8s} {t_match} | {em.fused_emotion:>8s} {f_match} | "
+ f"{em.fused_confidence:.3f}"
+ )
+
+ n = len(ground_truths)
+ logger.info(f"\nAccuracy: audio={audio_correct}/{n} text={text_correct}/{n} fused={fused_correct}/{n}")
+
+ # SpeakerSummary ์ถ๋ ฅ
+ logger.info("\nSpeaker Summaries:")
+ for spk, summary in output.speaker_summaries.items():
+ logger.info(f" {spk}: dominant={summary.dominant_emotion}, "
+ f"conf={summary.avg_confidence:.3f}, "
+ f"dist={summary.emotion_distribution}")
+
+ if errors:
+ logger.error(f"\nFAILED โ {len(errors)} errors:")
+ for e in errors:
+ logger.error(f" - {e}")
+ return False
+
+ logger.info(f"\n{'='*70}")
+ logger.info("E2E ํ
์คํธ PASS โ ์์ด ํ์ดํ๋ผ์ธ ์ ์ ๋์")
+ logger.info(f"{'='*70}")
+ return True
+
+
+if __name__ == "__main__":
+ success = run_e2e()
+ sys.exit(0 if success else 1)
diff --git a/scripts/test_meld_e2e_server.py b/scripts/test_meld_e2e_server.py
new file mode 100644
index 0000000000000000000000000000000000000000..5b0444545154de0e70010cf0546cf98ea15f7b3f
--- /dev/null
+++ b/scripts/test_meld_e2e_server.py
@@ -0,0 +1,253 @@
+#!/usr/bin/env python3
+"""MELD English test sets โ E2E server test.
+
+Uploads 8 MELD test WAVs to the deployed HF Spaces server,
+runs the full pipeline (Stage 1โ2โ3), and compares results
+against ground truth emotion labels.
+
+Usage:
+ python scripts/test_meld_e2e_server.py
+ python scripts/test_meld_e2e_server.py --server http://localhost:8000
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+import requests
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+MELD_DIR = PROJECT_ROOT / "data" / "meld_test"
+GT_PATH = MELD_DIR / "ground_truth.json"
+
+DEFAULT_SERVER = "https://bbbakery-ustwo-api.hf.space"
+POLL_INTERVAL = 5 # seconds
+MAX_WAIT = 300 # 5 minutes per file
+
+
+def health_check(base: str) -> bool:
+ try:
+ r = requests.get(f"{base}/api/health", timeout=10)
+ data = r.json()
+ if data.get("status") == "ok":
+ print(f" โ
Server OK ({data.get('timestamp', '?')})")
+ return True
+ except Exception as e:
+ print(f" โ Health check failed: {e}")
+ return False
+
+
+def upload(base: str, wav_path: Path) -> str | None:
+ """Upload WAV and return call_id."""
+ with open(wav_path, "rb") as f:
+ r = requests.post(
+ f"{base}/api/upload",
+ files={"file": (wav_path.name, f, "audio/wav")},
+ timeout=60,
+ )
+ if r.status_code != 200:
+ print(f" โ Upload failed ({r.status_code}): {r.text[:200]}")
+ return None
+ data = r.json()
+ return data.get("call_id")
+
+
+def analyze_and_poll(base: str, call_id: str) -> dict | None:
+ """Start analysis and poll until done."""
+ # Start
+ r = requests.post(f"{base}/api/analyze", params={"call_id": call_id}, timeout=30)
+ if r.status_code not in (200, 202):
+ print(f" โ Analyze start failed ({r.status_code}): {r.text[:200]}")
+ return None
+
+ data = r.json()
+ if data.get("status") == "done":
+ return data.get("result")
+
+ # Poll
+ elapsed = 0
+ while elapsed < MAX_WAIT:
+ time.sleep(POLL_INTERVAL)
+ elapsed += POLL_INTERVAL
+
+ r = requests.get(f"{base}/api/analyze/{call_id}/status", timeout=15)
+ data = r.json()
+ status = data.get("status")
+
+ if status == "done":
+ return data.get("result")
+ elif status == "error":
+ print(f" โ Pipeline error: {data.get('error', '?')}")
+ return None
+
+ mins, secs = divmod(elapsed, 60)
+ print(f" โณ {status}... ({int(mins)}m{int(secs)}s)", end="\r")
+
+ print(f" โ Timeout after {MAX_WAIT}s")
+ return None
+
+
+def extract_emotions(result: dict) -> dict:
+ """Extract emotion info from Stage 3 result."""
+ info = {}
+
+ # Character reactions โ emotions
+ reactions = result.get("character_reactions", [])
+ for i, rx in enumerate(reactions):
+ speaker = rx.get("speaker_id", f"speaker_{i}")
+ info[f"speaker_{i}_state"] = rx.get("solo_state", "?")
+
+ # Garden update
+ garden = result.get("garden_update", {})
+ info["garden_mood"] = garden.get("mood", "?")
+ info["garden_delta"] = garden.get("growth_delta", 0)
+
+ # Recap
+ recap = result.get("recap_card", {})
+ info["recap_headline"] = recap.get("headline", "?")
+
+ # Stage 2 emotions (if exposed in result)
+ stage2 = result.get("stage2_output", {})
+ if stage2:
+ for spk, summary in stage2.get("speaker_summaries", {}).items():
+ info[f"{spk}_dominant"] = summary.get("dominant_emotion", "?")
+ info[f"{spk}_distribution"] = summary.get("emotion_distribution", {})
+
+ return info
+
+
+def main():
+ parser = argparse.ArgumentParser(description="MELD E2E server test")
+ parser.add_argument("--server", default=DEFAULT_SERVER, help="Server base URL")
+ args = parser.parse_args()
+ base = args.server.rstrip("/")
+
+ print("=" * 70)
+ print(" MELD E2E Server Test")
+ print(f" Server: {base}")
+ print("=" * 70)
+
+ # Health check
+ print("\n[1] Health check")
+ if not health_check(base):
+ sys.exit(1)
+
+ # Load ground truth
+ print("\n[2] Loading ground truth")
+ with open(GT_PATH) as f:
+ gt = json.load(f)
+ print(f" {len(gt)} test sets loaded")
+
+ # Process each test set
+ print("\n[3] Running E2E tests\n")
+ results = {}
+ pass_count = 0
+ fail_count = 0
+
+ for tag in sorted(gt.keys()):
+ wav_path = MELD_DIR / f"{tag}.wav"
+ if not wav_path.exists():
+ print(f" โ ๏ธ {tag}: WAV not found, skipping")
+ continue
+
+ gt_entry = gt[tag]
+ print(f" ๐ฆ {tag} โ {gt_entry['description']}")
+ print(f" Primary: {gt_entry['primary_emotion']} | Duration: {gt_entry['duration_sec']}s | Utts: {gt_entry['total_utterances']}")
+
+ # Upload
+ call_id = upload(base, wav_path)
+ if not call_id:
+ fail_count += 1
+ results[tag] = {"status": "upload_failed"}
+ continue
+ print(f" Upload OK โ {call_id}")
+
+ # Analyze + poll
+ print(f" Analyzing...", end="")
+ start_time = time.time()
+ result = analyze_and_poll(base, call_id)
+ elapsed = time.time() - start_time
+
+ if not result:
+ fail_count += 1
+ results[tag] = {"status": "analyze_failed", "call_id": call_id}
+ print()
+ continue
+
+ print(f"\r โ
Done in {elapsed:.1f}s")
+
+ # Extract emotions
+ emotions = extract_emotions(result)
+
+ # Check pipeline completeness
+ has_reactions = len(result.get("character_reactions", [])) > 0
+ has_garden = "garden_update" in result
+ has_recap = "recap_card" in result
+
+ status = "pass" if (has_reactions and has_garden and has_recap) else "partial"
+ if status == "pass":
+ pass_count += 1
+ else:
+ fail_count += 1
+
+ results[tag] = {
+ "status": status,
+ "call_id": call_id,
+ "elapsed_sec": round(elapsed, 1),
+ "has_reactions": has_reactions,
+ "has_garden": has_garden,
+ "has_recap": has_recap,
+ "emotions": emotions,
+ "full_result": result,
+ "ground_truth": {
+ "primary_emotion": gt_entry["primary_emotion"],
+ "emotion_distribution": gt_entry["emotion_distribution"],
+ },
+ }
+
+ # Print details
+ print(f" Reactions: {'โ
' if has_reactions else 'โ'} | Garden: {'โ
' if has_garden else 'โ'} | Recap: {'โ
' if has_recap else 'โ'}")
+ for k, v in emotions.items():
+ if not k.startswith("full_"):
+ print(f" {k}: {v}")
+ print()
+
+ # Save results
+ out_path = MELD_DIR / "e2e_results.json"
+ with open(out_path, "w", encoding="utf-8") as f:
+ # Don't save full_result to keep file manageable
+ save_results = {}
+ for tag, r in results.items():
+ save_copy = {k: v for k, v in r.items() if k != "full_result"}
+ save_results[tag] = save_copy
+ json.dump(save_results, f, indent=2, ensure_ascii=False)
+
+ # Summary
+ print("=" * 70)
+ print(" SUMMARY")
+ print("=" * 70)
+ print(f"\n {'Tag':<25} {'Status':<10} {'Time':>6} {'Reactions':>10} {'Garden':>8} {'Recap':>7}")
+ print(" " + "-" * 70)
+ for tag in sorted(results.keys()):
+ r = results[tag]
+ status_icon = "โ
" if r["status"] == "pass" else "โ"
+ elapsed = f"{r.get('elapsed_sec', 0):.1f}s" if "elapsed_sec" in r else "โ"
+ react = "โ
" if r.get("has_reactions") else "โ"
+ garden = "โ
" if r.get("has_garden") else "โ"
+ recap = "โ
" if r.get("has_recap") else "โ"
+ print(f" {tag:<25} {status_icon:<10} {elapsed:>6} {react:>10} {garden:>8} {recap:>7}")
+
+ print(f"\n Total: {pass_count} pass / {fail_count} fail / {len(results)} total")
+ print(f" Results saved: {out_path}")
+ print("=" * 70)
+
+ return fail_count == 0
+
+
+if __name__ == "__main__":
+ success = main()
+ sys.exit(0 if success else 1)
diff --git a/scripts/train_emotion2vec.py b/scripts/train_emotion2vec.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/scripts/train_fusion_weights.py b/scripts/train_fusion_weights.py
new file mode 100644
index 0000000000000000000000000000000000000000..1577b116ce36efe77cf3e53d6274da05db257048
--- /dev/null
+++ b/scripts/train_fusion_weights.py
@@ -0,0 +1,297 @@
+#!/usr/bin/env python3
+"""Train per-emotion fusion weights via gradient descent.
+
+Inputs:
+ - --manifest JSON: list of {"path","text","label","source",...} (2,821 samples)
+ - --preds-cache JSON: {"audio_preds": [dict(7)], "text_preds": [dict(7)]}
+
+Output dir receives:
+ - trained_weights.json โ learned w_a / w_t / val_macro_f1
+ - trained_fusion_report.md โ comparison: audio-only, fixed 60/40, greedy optimal, trained
+ - trained_fusion_curve.png โ train/val loss + F1 curves
+
+Parameterization:
+ w_a[L] = sigmoid(ฮฑ[L]), w_t[L] = 1 - w_a[L] # 7 params total
+ fused[L] = p_a[L]*w_a[L] + p_t[L]*w_t[L]
+ fused โ normalize over L
+ loss = NLL(log fused, y) + ฮป * ||ฮฑ||ยฒ
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+from collections import Counter
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+from sklearn.metrics import f1_score
+from sklearn.model_selection import train_test_split
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+PROJECT_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear", "disgust"]
+
+
+class FusionHead(nn.Module):
+ def __init__(self, init_audio_frac: float = 0.6):
+ super().__init__()
+ init_val = float(torch.logit(torch.tensor(init_audio_frac)))
+ self.alpha = nn.Parameter(torch.full((7,), init_val))
+
+ @property
+ def w_a(self) -> torch.Tensor:
+ return torch.sigmoid(self.alpha)
+
+ @property
+ def w_t(self) -> torch.Tensor:
+ return 1.0 - self.w_a
+
+ def forward(self, p_a: torch.Tensor, p_t: torch.Tensor) -> torch.Tensor:
+ fused = p_a * self.w_a + p_t * self.w_t
+ return fused / fused.sum(dim=1, keepdim=True).clamp(min=1e-8)
+
+
+def probs_to_tensor(preds: list[dict]) -> torch.Tensor:
+ arr = np.array([[p.get(l, 0.0) for l in PROJECT_LABELS] for p in preds], dtype=np.float32)
+ return torch.from_numpy(arr)
+
+
+def map_label(lbl: str) -> str:
+ return "joy" if lbl == "happiness" else lbl
+
+
+def eval_weights(p_a: torch.Tensor, p_t: torch.Tensor, y: np.ndarray,
+ w_a_vec: np.ndarray) -> dict:
+ w_a = torch.from_numpy(w_a_vec.astype(np.float32))
+ w_t = 1.0 - w_a
+ fused = p_a * w_a + p_t * w_t
+ fused = fused / fused.sum(dim=1, keepdim=True).clamp(min=1e-8)
+ pred = fused.argmax(dim=1).numpy()
+ macro = f1_score(y, pred, average="macro")
+ per_class = {
+ PROJECT_LABELS[i]: f1_score((y == i).astype(int), (pred == i).astype(int))
+ for i in range(7)
+ }
+ return {"macro_f1": float(macro), "per_class": {k: float(v) for k, v in per_class.items()}}
+
+
+def train(p_a_tr, p_t_tr, y_tr, p_a_vl, p_t_vl, y_vl,
+ lr=0.05, epochs=500, l2=0.01, patience=50):
+ model = FusionHead()
+ opt = torch.optim.Adam(model.parameters(), lr=lr)
+ nll = nn.NLLLoss()
+
+ history = {"train_loss": [], "val_f1": []}
+ best_f1, best_alpha, waited = -1.0, None, 0
+
+ y_tr_t = torch.from_numpy(y_tr).long()
+ y_vl_t = torch.from_numpy(y_vl).long()
+
+ for epoch in range(epochs):
+ model.train()
+ opt.zero_grad()
+ fused = model(p_a_tr, p_t_tr)
+ loss = nll(torch.log(fused.clamp(min=1e-8)), y_tr_t) + l2 * (model.alpha ** 2).sum()
+ loss.backward()
+ opt.step()
+
+ model.eval()
+ with torch.no_grad():
+ val_fused = model(p_a_vl, p_t_vl)
+ val_pred = val_fused.argmax(dim=1).numpy()
+ val_f1 = f1_score(y_vl, val_pred, average="macro")
+
+ history["train_loss"].append(float(loss.item()))
+ history["val_f1"].append(float(val_f1))
+
+ if val_f1 > best_f1:
+ best_f1, best_alpha, waited = float(val_f1), model.alpha.detach().clone(), 0
+ else:
+ waited += 1
+ if waited >= patience:
+ logger.info("Early stop at epoch %d (patience=%d)", epoch, patience)
+ break
+
+ final_alpha = model.alpha.detach().clone()
+ model.alpha.data = best_alpha
+ return model, best_f1, history, final_alpha
+
+
+def plot_curve(history, output_path: Path) -> None:
+ import matplotlib
+ matplotlib.use("Agg")
+ import matplotlib.pyplot as plt
+
+ fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
+ ax1.plot(history["train_loss"], color="#F44336", label="train CE + L2")
+ ax1.set_xlabel("Epoch"); ax1.set_ylabel("Loss"); ax1.set_title("Training loss")
+ ax1.grid(alpha=0.3); ax1.legend()
+
+ ax2.plot(history["val_f1"], color="#4CAF50", label="val macro F1")
+ ax2.set_xlabel("Epoch"); ax2.set_ylabel("Macro F1"); ax2.set_title("Validation macro F1")
+ ax2.grid(alpha=0.3); ax2.legend()
+
+ plt.tight_layout()
+ plt.savefig(str(output_path), dpi=150)
+ plt.close()
+
+
+def write_report(output_path: Path, num_train: int, num_val: int, labels_dist: dict,
+ audio_only: dict, fixed: dict, greedy: dict, trained: dict,
+ trained_weights: dict, greedy_weights: dict) -> None:
+ lines = ["# Fusion Weight Training Report (v2)\n"]
+ lines.append(f"## Dataset\n\nTotal samples: **{num_train + num_val}** (train {num_train}, val {num_val})\n")
+ lines.append("### Label distribution\n\n| Label | Count |\n|---|---|")
+ for lbl, c in sorted(labels_dist.items(), key=lambda x: -x[1]):
+ lines.append(f"| {lbl} | {c} |")
+ lines.append("")
+ lines.append("## Macro F1 Comparison (validation set)\n")
+ lines.append("| Strategy | Macro F1 |")
+ lines.append("|---|---|")
+ lines.append(f"| Audio-only (argmax p_audio) | {audio_only['macro_f1']:.4f} |")
+ lines.append(f"| Fixed 60/40 | {fixed['macro_f1']:.4f} |")
+ lines.append(f"| Greedy grid (v1 weights) | {greedy['macro_f1']:.4f} |")
+ lines.append(f"| **Trained (gradient descent)** | **{trained['macro_f1']:.4f}** |")
+ lines.append("")
+ lines.append("## Per-class F1 (validation set)\n")
+ lines.append("| Emotion | Audio-only | Fixed 60/40 | Greedy | Trained |")
+ lines.append("|---|---|---|---|---|")
+ for lbl in PROJECT_LABELS:
+ lines.append(f"| {lbl} | {audio_only['per_class'][lbl]:.4f} | "
+ f"{fixed['per_class'][lbl]:.4f} | {greedy['per_class'][lbl]:.4f} | "
+ f"{trained['per_class'][lbl]:.4f} |")
+ lines.append("")
+ lines.append("## Learned weights\n")
+ lines.append("| Emotion | Audio (trained) | Text (trained) | Audio (greedy v1) |")
+ lines.append("|---|---|---|---|")
+ for lbl in PROJECT_LABELS:
+ w = trained_weights[lbl]
+ gw = greedy_weights.get(lbl, {"audio": None})
+ g_a = f"{gw['audio']:.2f}" if gw.get("audio") is not None else "โ"
+ lines.append(f"| {lbl} | {w['audio']:.2f} | {w['text']:.2f} | {g_a} |")
+ lines.append("")
+ output_path.write_text("\n".join(lines))
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--manifest", type=Path, required=True)
+ parser.add_argument("--preds-cache", type=Path, required=True)
+ parser.add_argument("--output-dir", type=Path, required=True)
+ parser.add_argument("--lr", type=float, default=0.05)
+ parser.add_argument("--epochs", type=int, default=500)
+ parser.add_argument("--l2", type=float, default=0.01)
+ parser.add_argument("--patience", type=int, default=50)
+ parser.add_argument("--seed", type=int, default=42)
+ parser.add_argument("--use-last-alpha", action="store_true",
+ help="Use final-epoch alpha instead of best-val-F1 alpha (keeps differentiation)")
+ args = parser.parse_args()
+
+ args.output_dir.mkdir(parents=True, exist_ok=True)
+
+ manifest = json.loads(args.manifest.read_text())
+ cache = json.loads(args.preds_cache.read_text())
+ audio_preds = cache["audio_preds"]
+ text_preds = cache["text_preds"]
+
+ if not (len(manifest) == len(audio_preds) == len(text_preds)):
+ raise ValueError(f"Size mismatch: manifest={len(manifest)}, audio={len(audio_preds)}, text={len(text_preds)}")
+
+ labels = [map_label(r["label"]) for r in manifest]
+ labels_dist = dict(Counter(labels))
+ logger.info("Loaded %d samples. Distribution: %s", len(manifest), labels_dist)
+
+ p_a = probs_to_tensor(audio_preds)
+ p_t = probs_to_tensor(text_preds)
+ y = np.array([PROJECT_LABELS.index(l) for l in labels])
+
+ # Stratified 80/20
+ tr_idx, vl_idx = train_test_split(
+ np.arange(len(y)), test_size=0.2, stratify=y, random_state=args.seed,
+ )
+ logger.info("Train/Val: %d / %d", len(tr_idx), len(vl_idx))
+
+ torch.manual_seed(args.seed)
+ model, best_f1, history, final_alpha = train(
+ p_a[tr_idx], p_t[tr_idx], y[tr_idx],
+ p_a[vl_idx], p_t[vl_idx], y[vl_idx],
+ lr=args.lr, epochs=args.epochs, l2=args.l2, patience=args.patience,
+ )
+ logger.info("Best val macro F1: %.4f", best_f1)
+
+ # Select which alpha to deploy: "best" (peak val F1) or "last" (final epoch)
+ if args.use_last_alpha:
+ model.alpha.data = final_alpha
+ logger.info("Using LAST-epoch alpha (--use-last-alpha)")
+ # Derive trained weights dict
+ w_a_np = model.w_a.detach().numpy()
+ trained_weights = {
+ PROJECT_LABELS[i]: {"audio": round(float(w_a_np[i]), 2), "text": round(float(1 - w_a_np[i]), 2)}
+ for i in range(7)
+ }
+
+ # Also compute last-alpha weights for comparison
+ last_w_a = torch.sigmoid(final_alpha).numpy()
+ last_weights = {
+ PROJECT_LABELS[i]: {"audio": round(float(last_w_a[i]), 2), "text": round(float(1 - last_w_a[i]), 2)}
+ for i in range(7)
+ }
+
+ # Baselines on val split
+ pa_vl = p_a[vl_idx]; pt_vl = p_t[vl_idx]; y_vl = y[vl_idx]
+
+ # Audio-only argmax
+ audio_only_pred = pa_vl.argmax(dim=1).numpy()
+ audio_only = {
+ "macro_f1": float(f1_score(y_vl, audio_only_pred, average="macro")),
+ "per_class": {
+ PROJECT_LABELS[i]: float(f1_score((y_vl == i).astype(int), (audio_only_pred == i).astype(int)))
+ for i in range(7)
+ },
+ }
+
+ fixed = eval_weights(pa_vl, pt_vl, y_vl, np.full(7, 0.6))
+ # Greedy v1 weights โ hardcode from previous report
+ GREEDY_V1 = {
+ "neutral": 0.75, "joy": 0.55, "sadness": 0.40, "anger": 0.65,
+ "surprise": 0.45, "fear": 0.00, "disgust": 0.80,
+ }
+ greedy_w_a = np.array([GREEDY_V1[l] for l in PROJECT_LABELS])
+ greedy = eval_weights(pa_vl, pt_vl, y_vl, greedy_w_a)
+ greedy_weights = {l: {"audio": GREEDY_V1[l]} for l in PROJECT_LABELS}
+
+ trained = eval_weights(pa_vl, pt_vl, y_vl, w_a_np)
+
+ logger.info("Audio-only val macro F1: %.4f", audio_only["macro_f1"])
+ logger.info("Fixed 60/40 val macro F1: %.4f", fixed["macro_f1"])
+ logger.info("Greedy v1 val macro F1: %.4f", greedy["macro_f1"])
+ logger.info("Trained val macro F1: %.4f", trained["macro_f1"])
+
+ # Save
+ (args.output_dir / "trained_weights.json").write_text(json.dumps({
+ "val_macro_f1": trained["macro_f1"],
+ "weights": trained_weights,
+ "last_epoch_weights": last_weights,
+ "train_size": int(len(tr_idx)),
+ "val_size": int(len(vl_idx)),
+ "hyperparams": {"lr": args.lr, "l2": args.l2, "epochs": args.epochs,
+ "patience": args.patience, "seed": args.seed,
+ "use_last_alpha": args.use_last_alpha},
+ }, indent=2, ensure_ascii=False))
+
+ plot_curve(history, args.output_dir / "trained_fusion_curve.png")
+ write_report(
+ args.output_dir / "trained_fusion_report.md",
+ len(tr_idx), len(vl_idx), labels_dist,
+ audio_only, fixed, greedy, trained,
+ trained_weights, greedy_weights,
+ )
+ logger.info("Done. Results saved to %s", args.output_dir)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/train_kcelectra.py b/scripts/train_kcelectra.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/scripts/train_lora_emotion2vec.py b/scripts/train_lora_emotion2vec.py
new file mode 100644
index 0000000000000000000000000000000000000000..910eb5e761dfdb0e5e5983877b1c9ac51267b5ce
--- /dev/null
+++ b/scripts/train_lora_emotion2vec.py
@@ -0,0 +1,749 @@
+#!/usr/bin/env python3
+"""Manual LoRA fine-tuning for emotion2vec_plus_base โ 7-class emotion.
+
+Wraps frozen attention layers with low-rank adapters (LoRALinear),
+replaces the 9-class proj head with a 7-class MLPHead, and trains
+with FocalLoss + disgust F1 gating.
+
+Usage:
+ # Quick smoke test (CPU)
+ python scripts/train_lora_emotion2vec.py \
+ --train-manifest data/lora_7class/train_manifest.json \
+ --val-manifest data/lora_7class/val_manifest.json \
+ --output-dir data/models/lora_emotion2vec_7class \
+ --epochs 3 --device cpu
+
+ # Full training (GPU, RTX 5050 8GB)
+ PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+ python scripts/train_lora_emotion2vec.py \
+ --train-manifest data/lora_7class/train_manifest.json \
+ --val-manifest data/lora_7class/val_manifest.json \
+ --output-dir data/models/lora_emotion2vec_7class \
+ --epochs 20 --batch-size 4 --accumulate-steps 8 --device cuda
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import random
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+# 7-class label taxonomy (matches prepare_lora_dataset.py)
+LABELS_7CLASS = ["happiness", "anger", "disgust", "fear", "neutral", "sadness", "surprise"]
+LABEL2IDX = {label: i for i, label in enumerate(LABELS_7CLASS)}
+NUM_CLASSES = len(LABELS_7CLASS)
+DISGUST_IDX = LABEL2IDX["disgust"] # 2
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# LoRA Components
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+class LoRALinear(nn.Module):
+ """Low-rank adapter wrapping a frozen nn.Linear.
+
+ At init, B is zero so LoRA contribution is zero (original behavior preserved).
+ scaling = alpha / r controls the magnitude of the LoRA update.
+ """
+
+ def __init__(self, original: nn.Linear, r: int = 16, alpha: int = 32, dropout: float = 0.1):
+ super().__init__()
+ self.original = original
+ self.r = r
+ self.scaling = alpha / r
+
+ # Freeze original weights
+ self.original.weight.requires_grad = False
+ if self.original.bias is not None:
+ self.original.bias.requires_grad = False
+
+ in_features = original.in_features
+ out_features = original.out_features
+
+ # Low-rank matrices
+ self.lora_A = nn.Linear(in_features, r, bias=False)
+ self.lora_B = nn.Linear(r, out_features, bias=False)
+ self.dropout = nn.Dropout(dropout)
+
+ # Init: A = kaiming, B = zero (so initial LoRA output = 0)
+ nn.init.kaiming_uniform_(self.lora_A.weight)
+ nn.init.zeros_(self.lora_B.weight)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ base_out = self.original(x)
+ lora_out = self.lora_B(self.dropout(self.lora_A(x))) * self.scaling
+ return base_out + lora_out
+
+
+def merge_lora_linear(lora: LoRALinear) -> nn.Linear:
+ """Merge LoRA weights into a plain nn.Linear for inference.
+
+ W_merged = W_original + scaling * B.weight @ A.weight
+ """
+ with torch.no_grad():
+ merged_weight = (
+ lora.original.weight
+ + lora.scaling * lora.lora_B.weight @ lora.lora_A.weight
+ )
+ bias = lora.original.bias
+
+ merged = nn.Linear(
+ lora.original.in_features,
+ lora.original.out_features,
+ bias=bias is not None,
+ )
+ merged.weight = nn.Parameter(merged_weight)
+ if bias is not None:
+ merged.bias = nn.Parameter(bias.clone())
+ return merged
+
+
+def inject_lora(
+ encoder: nn.Module,
+ r: int = 16,
+ alpha: int = 32,
+ dropout: float = 0.1,
+) -> None:
+ """Replace attn.qkv and attn.proj in each block with LoRALinear.
+
+ emotion2vec uses FUSED QKV: attn.qkv: Linear(768, 2304)
+ and attn.proj: Linear(768, 768). 8 blocks total.
+ MLP layers are NOT wrapped.
+ """
+ for block in encoder.blocks:
+ block.attn.qkv = LoRALinear(block.attn.qkv, r=r, alpha=alpha, dropout=dropout)
+ block.attn.proj = LoRALinear(block.attn.proj, r=r, alpha=alpha, dropout=dropout)
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Model Components
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+class MLPHead(nn.Module):
+ """Multi-layer classification head: 768 โ 512 โ 256 โ num_classes."""
+
+ def __init__(self, in_dim: int = 768, num_classes: int = NUM_CLASSES, dropout: float = 0.3):
+ super().__init__()
+ self.net = nn.Sequential(
+ nn.Linear(in_dim, 512),
+ nn.BatchNorm1d(512),
+ nn.GELU(),
+ nn.Dropout(dropout),
+ nn.Linear(512, 256),
+ nn.BatchNorm1d(256),
+ nn.GELU(),
+ nn.Dropout(dropout),
+ nn.Linear(256, num_classes),
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.net(x)
+
+
+class FocalLoss(nn.Module):
+ """Focal Loss with optional label smoothing and class weights."""
+
+ def __init__(self, weight=None, gamma: float = 2.0, label_smoothing: float = 0.05):
+ super().__init__()
+ self.gamma = gamma
+ self.weight = weight
+ self.label_smoothing = label_smoothing
+
+ def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+ ce_loss = F.cross_entropy(
+ logits, targets, weight=self.weight,
+ label_smoothing=self.label_smoothing, reduction="none",
+ )
+ pt = torch.exp(-ce_loss)
+ focal_loss = ((1 - pt) ** self.gamma) * ce_loss
+ return focal_loss.mean()
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Dataset
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+class EmotionDataset(Dataset):
+ """Load audio from unified manifest JSON for 7-class LoRA training.
+
+ Manifest format: list of {"audio_path": str, "label": str, ...}
+ OR {"path": str, "label": str, ...} (for backward compat).
+ """
+
+ def __init__(
+ self,
+ manifest_path: str,
+ max_duration_sec: float = 8.0,
+ phone_augment_prob: float = 0.0,
+ noise_augment_prob: float = 0.0,
+ ):
+ import torchaudio # lazy import
+
+ with open(manifest_path, encoding="utf-8") as f:
+ self.samples = json.load(f)
+
+ self.max_samples = int(max_duration_sec * 16000)
+ self.phone_augment_prob = phone_augment_prob
+ self.noise_augment_prob = noise_augment_prob
+ self._phone_sim = None
+
+ def _get_phone_sim(self):
+ if self._phone_sim is None:
+ sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+ from common.phone_simulator import PhoneSimulator, CompandingType
+ self._phone_sim = PhoneSimulator(companding=CompandingType.ALAW)
+ return self._phone_sim
+
+ def __len__(self):
+ return len(self.samples)
+
+ def __getitem__(self, idx):
+ import torchaudio
+
+ sample = self.samples[idx]
+ # Support both "audio_path" and "path" keys
+ audio_path = sample.get("audio_path") or sample.get("path", "")
+
+ waveform, sr = torchaudio.load(audio_path)
+ # Mono
+ if waveform.shape[0] > 1:
+ waveform = waveform.mean(dim=0, keepdim=True)
+ waveform = waveform.squeeze(0) # (T,)
+
+ # Resample to 16kHz
+ if sr != 16000:
+ waveform = torchaudio.functional.resample(waveform, sr, 16000)
+
+ # Truncate
+ if waveform.shape[0] > self.max_samples:
+ waveform = waveform[:self.max_samples]
+
+ audio = waveform.numpy()
+
+ # Augmentation
+ r = random.random()
+ if r < self.phone_augment_prob:
+ sim = self._get_phone_sim()
+ audio, _ = sim.process(audio, 16000)
+ import librosa
+ audio = librosa.resample(audio, orig_sr=8000, target_sr=16000)
+ elif r < self.phone_augment_prob + self.noise_augment_prob:
+ snr_db = random.uniform(10, 20)
+ signal_power = np.mean(audio ** 2)
+ noise_power = signal_power / (10 ** (snr_db / 10))
+ noise = np.random.normal(0, np.sqrt(max(noise_power, 1e-10)), len(audio)).astype(np.float32)
+ audio = audio + noise
+
+ label = LABEL2IDX[sample["label"]]
+ return torch.tensor(audio, dtype=torch.float32), label
+
+
+def collate_fn(batch):
+ """Pad waveforms to same length in batch."""
+ waveforms, labels = zip(*batch)
+ max_len = max(w.shape[0] for w in waveforms)
+ padded = torch.zeros(len(waveforms), max_len)
+ for i, w in enumerate(waveforms):
+ padded[i, :w.shape[0]] = w
+ return padded, torch.tensor(labels, dtype=torch.long)
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Model Loading & Forward
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def load_model(device: str, r: int = 16, alpha: int = 32, dropout: float = 0.1):
+ """Load emotion2vec_plus_base, freeze all, inject LoRA, replace proj."""
+ from funasr import AutoModel
+
+ logger.info("Loading emotion2vec_plus_base...")
+ fmodel = AutoModel(model="iic/emotion2vec_plus_base", device=device, hub="hf")
+ encoder = fmodel.model
+
+ # Freeze everything
+ for param in encoder.parameters():
+ param.requires_grad = False
+
+ # Inject LoRA adapters into attention layers
+ inject_lora(encoder, r=r, alpha=alpha, dropout=dropout)
+
+ # Replace 9-class proj with 7-class MLPHead
+ old_proj = encoder.proj
+ encoder.proj = MLPHead(768, NUM_CLASSES)
+ logger.info("Replaced proj: Linear(768, %d) -> MLPHead(768->512->256->%d)",
+ old_proj.out_features, NUM_CLASSES)
+
+ # Move entire model (including new LoRA params + MLPHead) to device
+ encoder = encoder.to(device)
+ return encoder
+
+
+def forward_pass(encoder, waveforms: torch.Tensor, device: str) -> torch.Tensor:
+ """Differentiable forward pass through emotion2vec with LoRA.
+
+ Args:
+ encoder: emotion2vec model with LoRA injected
+ waveforms: (B, T) float32, 16kHz
+ device: compute device
+
+ Returns:
+ logits: (B, 7)
+ """
+ waveforms = waveforms.to(device)
+
+ # Layer norm (per emotion2vec inference)
+ if encoder.cfg.normalize:
+ normed = []
+ for i in range(waveforms.shape[0]):
+ normed.append(F.layer_norm(waveforms[i], waveforms[i].shape))
+ waveforms = torch.stack(normed)
+
+ # Extract features
+ feats = encoder.extract_features(waveforms, padding_mask=None)
+ x = feats["x"] # (B, T', 768)
+
+ # Mean pool + classify
+ pooled = x.mean(dim=1) # (B, 768)
+ logits = encoder.proj(pooled) # (B, 7)
+ return logits
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Validation
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+@torch.no_grad()
+def validate(encoder, val_loader, device, criterion):
+ """Run validation, return metrics including disgust_f1."""
+ encoder.eval()
+ total_loss = 0
+ y_true, y_pred = [], []
+
+ for waveforms, labels in val_loader:
+ labels = labels.to(device)
+ logits = forward_pass(encoder, waveforms, device)
+ loss = criterion(logits, labels)
+ total_loss += loss.item() * labels.size(0)
+
+ preds = logits.argmax(dim=-1)
+ y_true.extend(labels.cpu().tolist())
+ y_pred.extend(preds.cpu().tolist())
+
+ from sklearn.metrics import precision_recall_fscore_support, accuracy_score, confusion_matrix
+
+ accuracy = accuracy_score(y_true, y_pred)
+ _, _, f1_per_class, _ = precision_recall_fscore_support(
+ y_true, y_pred, labels=list(range(NUM_CLASSES)), average=None, zero_division=0,
+ )
+ macro_f1 = float(np.mean(f1_per_class))
+ disgust_f1 = float(f1_per_class[DISGUST_IDX])
+
+ per_class = {LABELS_7CLASS[i]: round(float(f1_per_class[i]), 4) for i in range(NUM_CLASSES)}
+ avg_loss = total_loss / max(len(y_true), 1)
+ cm = confusion_matrix(y_true, y_pred, labels=list(range(NUM_CLASSES)))
+
+ return {
+ "loss": round(avg_loss, 4),
+ "accuracy": round(accuracy, 4),
+ "macro_f1": round(macro_f1, 4),
+ "disgust_f1": round(disgust_f1, 4),
+ "per_class_f1": per_class,
+ "confusion_matrix": cm.tolist(),
+ "y_true": y_true,
+ "y_pred": y_pred,
+ }
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Checkpoint
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def save_lora_checkpoint(encoder, path: Path, epoch: int, metrics: dict,
+ best_f1: float = 0.0, patience_counter: int = 0,
+ optimizer=None, scheduler=None, scaler=None,
+ training_log=None):
+ """Save LoRA weights + MLPHead + optimizer state for resume."""
+ lora_state = {}
+ for name, module in encoder.named_modules():
+ if isinstance(module, LoRALinear):
+ lora_state[f"{name}.lora_A.weight"] = module.lora_A.weight.data.cpu()
+ lora_state[f"{name}.lora_B.weight"] = module.lora_B.weight.data.cpu()
+
+ state = {
+ "lora_weights": lora_state,
+ "proj": encoder.proj.state_dict(),
+ "epoch": epoch,
+ "metrics": metrics,
+ "best_f1": best_f1,
+ "patience_counter": patience_counter,
+ "num_classes": NUM_CLASSES,
+ "labels": LABELS_7CLASS,
+ }
+ if optimizer is not None:
+ state["optimizer"] = optimizer.state_dict()
+ if scheduler is not None:
+ state["scheduler"] = scheduler.state_dict()
+ if scaler is not None:
+ state["scaler"] = scaler.state_dict()
+ if training_log is not None:
+ state["training_log"] = training_log
+ torch.save(state, str(path))
+
+
+def load_lora_checkpoint(encoder, path: Path, device: str,
+ optimizer=None, scheduler=None, scaler=None):
+ """Load LoRA checkpoint and restore training state."""
+ logger.info("Resuming from checkpoint: %s", path)
+ ckpt = torch.load(str(path), map_location=device, weights_only=False)
+
+ # Restore LoRA weights
+ for name, module in encoder.named_modules():
+ if isinstance(module, LoRALinear):
+ a_key = f"{name}.lora_A.weight"
+ b_key = f"{name}.lora_B.weight"
+ if a_key in ckpt["lora_weights"]:
+ module.lora_A.weight.data.copy_(ckpt["lora_weights"][a_key].to(device))
+ module.lora_B.weight.data.copy_(ckpt["lora_weights"][b_key].to(device))
+
+ # Restore MLPHead
+ encoder.proj.load_state_dict(ckpt["proj"])
+
+ # Restore optimizer/scheduler/scaler if available
+ if optimizer is not None and "optimizer" in ckpt:
+ optimizer.load_state_dict(ckpt["optimizer"])
+ if scheduler is not None and "scheduler" in ckpt:
+ scheduler.load_state_dict(ckpt["scheduler"])
+ if scaler is not None and "scaler" in ckpt:
+ scaler.load_state_dict(ckpt["scaler"])
+
+ logger.info("Resumed from epoch %d (best_f1=%.4f, patience=%d)",
+ ckpt["epoch"], ckpt["best_f1"], ckpt["patience_counter"])
+ return {
+ "epoch": ckpt["epoch"],
+ "best_f1": ckpt["best_f1"],
+ "patience_counter": ckpt["patience_counter"],
+ "training_log": ckpt.get("training_log", []),
+ }
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Confusion Matrix Plot
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def plot_confusion_matrix(cm, output_path: Path, epoch: int):
+ """Save confusion matrix as PNG."""
+ try:
+ import matplotlib
+ matplotlib.use("Agg")
+ import matplotlib.pyplot as plt
+ import seaborn as sns
+
+ fig, ax = plt.subplots(figsize=(9, 7))
+ cm_norm = cm / np.maximum(cm.sum(axis=1, keepdims=True), 1)
+ sns.heatmap(
+ cm_norm, annot=True, fmt=".2f", cmap="Blues",
+ xticklabels=LABELS_7CLASS, yticklabels=LABELS_7CLASS, ax=ax,
+ )
+ for i in range(NUM_CLASSES):
+ for j in range(NUM_CLASSES):
+ ax.text(j + 0.5, i + 0.7, f"({cm[i][j]})",
+ ha="center", va="center", fontsize=6, color="gray")
+
+ ax.set_xlabel("Predicted")
+ ax.set_ylabel("True")
+ ax.set_title(f"LoRA 7-Class Confusion Matrix (Epoch {epoch})")
+ plt.tight_layout()
+ plt.savefig(str(output_path), dpi=150)
+ plt.close()
+ logger.info("Confusion matrix saved to %s", output_path)
+ except ImportError:
+ logger.warning("matplotlib/seaborn not available โ skipping confusion matrix plot")
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Training
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def train(args):
+ device = args.device
+ use_amp = (device == "cuda")
+ accumulate_steps = args.accumulate_steps
+
+ # Load model
+ encoder = load_model(device, r=args.lora_r, alpha=args.lora_alpha, dropout=args.lora_dropout)
+
+ # Log trainable vs total
+ trainable = sum(p.numel() for p in encoder.parameters() if p.requires_grad)
+ total = sum(p.numel() for p in encoder.parameters())
+ logger.info("Trainable: %dK / Total: %dK (%.1f%%)",
+ trainable // 1000, total // 1000, 100 * trainable / total)
+
+ # Datasets
+ train_dataset = EmotionDataset(
+ args.train_manifest,
+ max_duration_sec=8.0,
+ phone_augment_prob=args.phone_augment_prob,
+ noise_augment_prob=args.noise_augment_prob,
+ )
+ val_dataset = EmotionDataset(args.val_manifest, max_duration_sec=8.0)
+
+ logger.info("Train: %d samples, Val: %d samples", len(train_dataset), len(val_dataset))
+ logger.info("AMP: %s, Accumulate: %d, Effective batch: %d",
+ use_amp, accumulate_steps, args.batch_size * accumulate_steps)
+
+ n_workers = 2
+ use_pin = (device == "cuda")
+ train_loader = DataLoader(
+ train_dataset, batch_size=args.batch_size, shuffle=True,
+ collate_fn=collate_fn, num_workers=n_workers, drop_last=True,
+ pin_memory=use_pin, persistent_workers=(n_workers > 0),
+ )
+ val_loader = DataLoader(
+ val_dataset, batch_size=args.batch_size, shuffle=False,
+ collate_fn=collate_fn, num_workers=n_workers,
+ pin_memory=use_pin, persistent_workers=(n_workers > 0),
+ )
+
+ # Class weights: inverse frequency + disgust 2.5x boost
+ class_counts = np.zeros(NUM_CLASSES)
+ for sample in train_dataset.samples:
+ class_counts[LABEL2IDX[sample["label"]]] += 1
+ class_weights = 1.0 / np.maximum(class_counts, 1)
+ class_weights = class_weights / class_weights.sum() * NUM_CLASSES
+ class_weights[DISGUST_IDX] *= 2.5 # Disgust boost
+ logger.info("Class weights: %s",
+ {LABELS_7CLASS[i]: round(float(class_weights[i]), 3) for i in range(NUM_CLASSES)})
+
+ criterion = FocalLoss(
+ weight=torch.tensor(class_weights, dtype=torch.float32).to(device),
+ gamma=2.0,
+ label_smoothing=0.05,
+ )
+
+ # Optimizer: differential LR
+ lora_params = []
+ proj_params = list(encoder.proj.parameters())
+ proj_ids = {id(p) for p in proj_params}
+ for name, param in encoder.named_parameters():
+ if param.requires_grad and id(param) not in proj_ids:
+ lora_params.append(param)
+
+ optimizer = torch.optim.AdamW([
+ {"params": lora_params, "lr": args.lora_lr},
+ {"params": proj_params, "lr": args.proj_lr},
+ ], weight_decay=args.weight_decay)
+
+ # OneCycleLR scheduler
+ steps_per_epoch = max(len(train_loader) // accumulate_steps, 1)
+ total_steps = steps_per_epoch * args.epochs
+ scheduler = torch.optim.lr_scheduler.OneCycleLR(
+ optimizer,
+ max_lr=[args.lora_lr, args.proj_lr],
+ total_steps=total_steps,
+ pct_start=0.1,
+ anneal_strategy="cos",
+ )
+
+ # Output dir
+ output_dir = Path(args.output_dir)
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ # AMP scaler
+ scaler = torch.amp.GradScaler("cuda", enabled=use_amp)
+
+ # Training state
+ training_log = []
+ best_f1 = 0.0
+ patience_counter = 0
+ start_epoch = 1
+ disgust_gate = 0.3
+
+ # Resume from checkpoint if requested
+ if args.resume:
+ resume_path = output_dir / "last_lora.pt"
+ if resume_path.exists():
+ resume_state = load_lora_checkpoint(
+ encoder, resume_path, device,
+ optimizer=optimizer, scheduler=scheduler, scaler=scaler,
+ )
+ start_epoch = resume_state["epoch"] + 1
+ best_f1 = resume_state["best_f1"]
+ patience_counter = resume_state["patience_counter"]
+ training_log = resume_state["training_log"]
+ logger.info("Resuming training from epoch %d (best_f1=%.4f)", start_epoch, best_f1)
+ else:
+ logger.warning("--resume set but no checkpoint found at %s, starting fresh", resume_path)
+
+ for epoch in range(start_epoch, args.epochs + 1):
+ epoch_start = time.time()
+ encoder.train()
+ total_loss = 0
+ correct = 0
+ total_samples = 0
+ optimizer.zero_grad()
+
+ for batch_idx, (waveforms, labels) in enumerate(train_loader):
+ labels = labels.to(device)
+
+ with torch.amp.autocast("cuda", enabled=use_amp):
+ logits = forward_pass(encoder, waveforms, device)
+ loss = criterion(logits, labels) / accumulate_steps
+
+ scaler.scale(loss).backward()
+
+ if (batch_idx + 1) % accumulate_steps == 0 or (batch_idx + 1) == len(train_loader):
+ scaler.unscale_(optimizer)
+ trainable_params = [p for p in encoder.parameters() if p.requires_grad]
+ torch.nn.utils.clip_grad_norm_(trainable_params, 1.0)
+ scaler.step(optimizer)
+ scaler.update()
+ optimizer.zero_grad()
+ scheduler.step()
+
+ total_loss += loss.item() * accumulate_steps * labels.size(0)
+ preds = logits.argmax(dim=-1)
+ correct += (preds == labels).sum().item()
+ total_samples += labels.size(0)
+
+ if (batch_idx + 1) % 10 == 0:
+ cur_lr = optimizer.param_groups[0]["lr"]
+ logger.info(" Epoch %d [%d/%d] loss=%.4f lr=%.2e",
+ epoch, batch_idx + 1, len(train_loader),
+ loss.item() * accumulate_steps, cur_lr)
+
+ train_loss = total_loss / max(total_samples, 1)
+ train_acc = correct / max(total_samples, 1)
+
+ # Validate
+ val_metrics = validate(encoder, val_loader, device, criterion)
+
+ epoch_time = time.time() - epoch_start
+ logger.info(
+ "Epoch %d/%d (%.0fs): train_loss=%.4f train_acc=%.3f | "
+ "val_loss=%.4f val_f1=%.3f val_acc=%.3f disgust_f1=%.3f",
+ epoch, args.epochs, epoch_time,
+ train_loss, train_acc,
+ val_metrics["loss"], val_metrics["macro_f1"],
+ val_metrics["accuracy"], val_metrics["disgust_f1"],
+ )
+ logger.info(" Per-class F1: %s", val_metrics["per_class_f1"])
+
+ epoch_log = {
+ "epoch": epoch,
+ "train_loss": round(train_loss, 4),
+ "train_acc": round(train_acc, 4),
+ "val_loss": val_metrics["loss"],
+ "val_accuracy": val_metrics["accuracy"],
+ "val_macro_f1": val_metrics["macro_f1"],
+ "val_disgust_f1": val_metrics["disgust_f1"],
+ "val_per_class_f1": val_metrics["per_class_f1"],
+ "epoch_time_sec": round(epoch_time, 1),
+ }
+ training_log.append(epoch_log)
+
+ # Disgust F1 Gate + best model save
+ gate_pass = val_metrics["disgust_f1"] >= disgust_gate
+ if not gate_pass:
+ logger.warning("[GATE FAIL] disgust_f1=%.3f < %.1f โ not saving as best",
+ val_metrics["disgust_f1"], disgust_gate)
+
+ if val_metrics["macro_f1"] > best_f1 and gate_pass:
+ best_f1 = val_metrics["macro_f1"]
+ patience_counter = 0
+ save_lora_checkpoint(
+ encoder, output_dir / "best_lora.pt", epoch, val_metrics, best_f1,
+ optimizer=optimizer, scheduler=scheduler, scaler=scaler,
+ training_log=training_log,
+ )
+ logger.info(" New best! macro_f1=%.4f (gate passed) saved to best_lora.pt", best_f1)
+
+ if "confusion_matrix" in val_metrics:
+ cm = np.array(val_metrics["confusion_matrix"])
+ plot_confusion_matrix(cm, output_dir / f"confusion_matrix_epoch{epoch}.png", epoch)
+ elif gate_pass:
+ patience_counter += 1
+ else:
+ # Gate fail does not increment patience
+ pass
+
+ # Always save last (with full state for resume)
+ save_lora_checkpoint(
+ encoder, output_dir / "last_lora.pt", epoch, val_metrics, best_f1, patience_counter,
+ optimizer=optimizer, scheduler=scheduler, scaler=scaler,
+ training_log=training_log,
+ )
+
+ # Save training log
+ with open(output_dir / "training_log.json", "w") as f:
+ json.dump(training_log, f, indent=2)
+
+ # Early stopping (only on gate-passing epochs)
+ if patience_counter >= args.patience:
+ logger.info("Early stopping at epoch %d (patience=%d)", epoch, args.patience)
+ break
+
+ if device == "cuda":
+ torch.cuda.empty_cache()
+
+ # Save config
+ config = {
+ "model": "iic/emotion2vec_plus_base",
+ "method": "LoRA",
+ "lora_r": args.lora_r,
+ "lora_alpha": args.lora_alpha,
+ "num_classes": NUM_CLASSES,
+ "labels": LABELS_7CLASS,
+ "label2idx": LABEL2IDX,
+ "best_val_f1": best_f1,
+ "training_args": vars(args),
+ }
+ with open(output_dir / "config.json", "w") as f:
+ json.dump(config, f, indent=2, ensure_ascii=False)
+
+ logger.info("Training complete. Best F1=%.4f at %s", best_f1, output_dir / "best_lora.pt")
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Main
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def main():
+ parser = argparse.ArgumentParser(description="LoRA fine-tune emotion2vec 7-class")
+ parser.add_argument("--train-manifest", required=True)
+ parser.add_argument("--val-manifest", required=True)
+ parser.add_argument("--output-dir", default="data/models/lora_emotion2vec_7class")
+ parser.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
+ parser.add_argument("--epochs", type=int, default=20)
+ parser.add_argument("--batch-size", type=int, default=4)
+ parser.add_argument("--accumulate-steps", type=int, default=8)
+ parser.add_argument("--lora-r", type=int, default=16)
+ parser.add_argument("--lora-alpha", type=int, default=32)
+ parser.add_argument("--lora-dropout", type=float, default=0.1)
+ parser.add_argument("--lora-lr", type=float, default=2e-4)
+ parser.add_argument("--proj-lr", type=float, default=2e-3)
+ parser.add_argument("--weight-decay", type=float, default=0.01)
+ parser.add_argument("--patience", type=int, default=7)
+ parser.add_argument("--phone-augment-prob", type=float, default=0.3)
+ parser.add_argument("--noise-augment-prob", type=float, default=0.15)
+ parser.add_argument("--resume", action="store_true",
+ help="Resume from last_lora.pt checkpoint in output-dir")
+ args = parser.parse_args()
+
+ train(args)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/train_lora_kcelectra.py b/scripts/train_lora_kcelectra.py
new file mode 100644
index 0000000000000000000000000000000000000000..4d6082a8a0d3f4c0f980a873e0dfd36f3e88a3b2
--- /dev/null
+++ b/scripts/train_lora_kcelectra.py
@@ -0,0 +1,402 @@
+#!/usr/bin/env python3
+"""LoRA Fine-Tuning KcELECTRA for 7-Class Korean Text Emotion Recognition.
+
+Uses PEFT LoRA on beomi/KcELECTRA-base-v2022 for text-based emotion classification.
+Filters out samples without text (e.g., RAVDESS English data).
+
+Usage:
+ python scripts/train_lora_kcelectra.py \
+ --train-manifest data/lora_dataset/train_manifest.json \
+ --val-manifest data/lora_dataset/val_manifest.json \
+ --output-dir data/models/lora_kcelectra_7class \
+ --epochs 10 --batch-size 16 --device cuda
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import random
+import time
+from collections import Counter
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, Dataset
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+LABELS_7CLASS = ["happiness", "anger", "disgust", "fear", "neutral", "sadness", "surprise"]
+LABEL2IDX = {l: i for i, l in enumerate(LABELS_7CLASS)}
+NUM_CLASSES = len(LABELS_7CLASS)
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Dataset
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+class TextEmotionDataset(Dataset):
+ """Load text + label from manifest, skip samples without text."""
+
+ def __init__(self, manifest_path: str, tokenizer, max_length: int = 128):
+ with open(manifest_path, encoding="utf-8") as f:
+ raw = json.load(f)
+
+ # Filter: only samples with non-empty Korean text
+ self.samples = [
+ s for s in raw
+ if s.get("text", "").strip() and s["label"] in LABEL2IDX
+ ]
+ self.tokenizer = tokenizer
+ self.max_length = max_length
+
+ logger.info("TextEmotionDataset: %d samples (filtered from %d, skipped %d without text)",
+ len(self.samples), len(raw), len(raw) - len(self.samples))
+
+ def __len__(self):
+ return len(self.samples)
+
+ def __getitem__(self, idx):
+ sample = self.samples[idx]
+ text = sample["text"]
+ label = LABEL2IDX[sample["label"]]
+
+ encoding = self.tokenizer(
+ text,
+ truncation=True,
+ max_length=self.max_length,
+ padding="max_length",
+ return_tensors="pt",
+ )
+
+ return {
+ "input_ids": encoding["input_ids"].squeeze(0),
+ "attention_mask": encoding["attention_mask"].squeeze(0),
+ "labels": torch.tensor(label, dtype=torch.long),
+ }
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Validation
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+@torch.no_grad()
+def validate(model, val_loader, device, criterion):
+ model.eval()
+ total_loss = 0
+ y_true, y_pred = [], []
+
+ for batch in val_loader:
+ input_ids = batch["input_ids"].to(device)
+ attention_mask = batch["attention_mask"].to(device)
+ labels = batch["labels"].to(device)
+
+ outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+ logits = outputs.logits
+ loss = criterion(logits, labels)
+
+ total_loss += loss.item() * labels.size(0)
+ y_true.extend(labels.cpu().tolist())
+ y_pred.extend(logits.argmax(dim=-1).cpu().tolist())
+
+ from sklearn.metrics import accuracy_score, f1_score, confusion_matrix
+ acc = accuracy_score(y_true, y_pred)
+ f1_per_class = f1_score(y_true, y_pred, labels=list(range(NUM_CLASSES)),
+ average=None, zero_division=0)
+ macro_f1 = float(np.mean(f1_per_class))
+ per_class = {LABELS_7CLASS[i]: round(float(f1_per_class[i]), 4) for i in range(NUM_CLASSES)}
+ cm = confusion_matrix(y_true, y_pred, labels=list(range(NUM_CLASSES)))
+
+ return {
+ "loss": round(total_loss / max(len(y_true), 1), 4),
+ "accuracy": round(acc, 4),
+ "macro_f1": round(macro_f1, 4),
+ "per_class_f1": per_class,
+ "confusion_matrix": cm.tolist(),
+ }
+
+
+def plot_confusion_matrix(cm, output_path: Path, epoch: int):
+ try:
+ import matplotlib; matplotlib.use("Agg")
+ import matplotlib.pyplot as plt
+ import seaborn as sns
+ fig, ax = plt.subplots(figsize=(9, 7))
+ cm_norm = cm / cm.sum(axis=1, keepdims=True)
+ sns.heatmap(cm_norm, annot=True, fmt=".2f", cmap="Blues",
+ xticklabels=LABELS_7CLASS, yticklabels=LABELS_7CLASS, ax=ax)
+ for i in range(NUM_CLASSES):
+ for j in range(NUM_CLASSES):
+ ax.text(j + 0.5, i + 0.7, f"({cm[i][j]})",
+ ha="center", va="center", fontsize=6, color="gray")
+ ax.set_xlabel("Predicted"); ax.set_ylabel("True")
+ ax.set_title(f"KcELECTRA LoRA โ Epoch {epoch}")
+ plt.tight_layout(); plt.savefig(str(output_path), dpi=150); plt.close()
+ except ImportError:
+ pass
+
+
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+# Training
+# โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def train(args):
+ device = args.device
+ use_amp = (device == "cuda")
+
+ # Load tokenizer + base model
+ from transformers import AutoTokenizer, AutoModelForSequenceClassification
+ from peft import LoraConfig, get_peft_model, TaskType
+
+ logger.info("Loading KcELECTRA: %s", args.model_id)
+ tokenizer = AutoTokenizer.from_pretrained(args.model_id)
+ model = AutoModelForSequenceClassification.from_pretrained(
+ args.model_id,
+ num_labels=NUM_CLASSES,
+ id2label={i: l for i, l in enumerate(LABELS_7CLASS)},
+ label2id=LABEL2IDX,
+ )
+
+ # Apply LoRA
+ lora_config = LoraConfig(
+ r=args.lora_r,
+ lora_alpha=args.lora_alpha,
+ lora_dropout=args.lora_dropout,
+ target_modules=["query", "value"],
+ task_type=TaskType.SEQ_CLS,
+ bias="none",
+ )
+ model = get_peft_model(model, lora_config)
+ model.print_trainable_parameters()
+ model = model.to(device)
+
+ # Datasets
+ train_ds = TextEmotionDataset(args.train_manifest, tokenizer, max_length=args.max_length)
+ val_ds = TextEmotionDataset(args.val_manifest, tokenizer, max_length=args.max_length)
+ logger.info("Train: %d, Val: %d", len(train_ds), len(val_ds))
+
+ train_loader = DataLoader(
+ train_ds, batch_size=args.batch_size, shuffle=True,
+ num_workers=2, pin_memory=(device == "cuda"), drop_last=True,
+ )
+ val_loader = DataLoader(
+ val_ds, batch_size=args.batch_size, shuffle=False,
+ num_workers=2, pin_memory=(device == "cuda"),
+ )
+
+ # Class weights: inverse frequency only (no extra boost)
+ class_counts = np.zeros(NUM_CLASSES)
+ for s in train_ds.samples:
+ class_counts[LABEL2IDX[s["label"]]] += 1
+ class_weights = 1.0 / np.maximum(class_counts, 1)
+ class_weights = class_weights / class_weights.sum() * NUM_CLASSES
+ logger.info("Class weights: %s",
+ {LABELS_7CLASS[i]: round(float(class_weights[i]), 3) for i in range(NUM_CLASSES)})
+
+ criterion = nn.CrossEntropyLoss(
+ weight=torch.tensor(class_weights, dtype=torch.float32).to(device),
+ )
+
+ # Optimizer
+ optimizer = torch.optim.AdamW(
+ [p for p in model.parameters() if p.requires_grad],
+ lr=args.lr,
+ weight_decay=args.weight_decay,
+ )
+
+ # Scheduler
+ steps_per_epoch = max(len(train_loader) // args.accumulate_steps, 1)
+ total_steps = steps_per_epoch * args.epochs
+ scheduler = torch.optim.lr_scheduler.OneCycleLR(
+ optimizer, max_lr=args.lr, total_steps=total_steps,
+ pct_start=0.1, anneal_strategy="cos",
+ )
+
+ scaler = torch.amp.GradScaler("cuda", enabled=use_amp)
+ output_dir = Path(args.output_dir)
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ # Training state
+ training_log = []
+ best_f1 = 0.0
+ patience_counter = 0
+ start_epoch = 1
+
+ # Resume
+ if args.resume:
+ ckpt_path = output_dir / "last_checkpoint.json"
+ if ckpt_path.exists():
+ with open(ckpt_path) as f:
+ ckpt_info = json.load(f)
+ start_epoch = ckpt_info["epoch"] + 1
+ best_f1 = ckpt_info["best_f1"]
+ patience_counter = ckpt_info["patience_counter"]
+ training_log = ckpt_info.get("training_log", [])
+ # Load model weights
+ model_path = output_dir / "last_model"
+ if model_path.exists():
+ from peft import PeftModel
+ model = AutoModelForSequenceClassification.from_pretrained(
+ args.model_id, num_labels=NUM_CLASSES,
+ id2label={i: l for i, l in enumerate(LABELS_7CLASS)},
+ label2id=LABEL2IDX,
+ )
+ model = PeftModel.from_pretrained(model, str(model_path))
+ model = model.to(device)
+ # Rebuild optimizer
+ optimizer = torch.optim.AdamW(
+ [p for p in model.parameters() if p.requires_grad],
+ lr=args.lr, weight_decay=args.weight_decay,
+ )
+ remaining_steps = steps_per_epoch * (args.epochs - start_epoch + 1)
+ scheduler = torch.optim.lr_scheduler.OneCycleLR(
+ optimizer, max_lr=args.lr, total_steps=max(remaining_steps, 1),
+ pct_start=0.1, anneal_strategy="cos",
+ )
+ logger.info("Resumed from epoch %d (best_f1=%.4f)", start_epoch, best_f1)
+ else:
+ logger.warning("--resume but no checkpoint found, starting fresh")
+
+ for epoch in range(start_epoch, args.epochs + 1):
+ epoch_start = time.time()
+ model.train()
+ total_loss = 0; correct = 0; total_samples = 0
+ optimizer.zero_grad()
+
+ for batch_idx, batch in enumerate(train_loader):
+ input_ids = batch["input_ids"].to(device)
+ attention_mask = batch["attention_mask"].to(device)
+ labels = batch["labels"].to(device)
+
+ with torch.amp.autocast("cuda", enabled=use_amp):
+ outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+ logits = outputs.logits
+ loss = criterion(logits, labels) / args.accumulate_steps
+
+ scaler.scale(loss).backward()
+
+ if (batch_idx + 1) % args.accumulate_steps == 0 or (batch_idx + 1) == len(train_loader):
+ scaler.unscale_(optimizer)
+ torch.nn.utils.clip_grad_norm_(
+ [p for p in model.parameters() if p.requires_grad], 1.0,
+ )
+ scaler.step(optimizer); scaler.update()
+ optimizer.zero_grad(); scheduler.step()
+
+ total_loss += loss.item() * args.accumulate_steps * labels.size(0)
+ preds = logits.argmax(dim=-1)
+ correct += (preds == labels).sum().item()
+ total_samples += labels.size(0)
+
+ if (batch_idx + 1) % 50 == 0:
+ logger.info(" Epoch %d [%d/%d] loss=%.4f lr=%.2e",
+ epoch, batch_idx + 1, len(train_loader),
+ loss.item() * args.accumulate_steps,
+ optimizer.param_groups[0]["lr"])
+
+ train_loss = total_loss / max(total_samples, 1)
+ train_acc = correct / max(total_samples, 1)
+
+ # Validate
+ val_metrics = validate(model, val_loader, device, criterion)
+ epoch_time = time.time() - epoch_start
+
+ logger.info(
+ "Epoch %d/%d (%.0fs): train_loss=%.4f train_acc=%.3f | "
+ "val_loss=%.4f val_f1=%.3f val_acc=%.3f",
+ epoch, args.epochs, epoch_time, train_loss, train_acc,
+ val_metrics["loss"], val_metrics["macro_f1"], val_metrics["accuracy"],
+ )
+ logger.info(" Per-class F1: %s", val_metrics["per_class_f1"])
+
+ training_log.append({
+ "epoch": epoch,
+ "train_loss": round(train_loss, 4),
+ "train_acc": round(train_acc, 4),
+ **{f"val_{k}": v for k, v in val_metrics.items() if k != "confusion_matrix"},
+ "epoch_time_sec": round(epoch_time, 1),
+ })
+
+ # Best model
+ if val_metrics["macro_f1"] > best_f1:
+ best_f1 = val_metrics["macro_f1"]
+ patience_counter = 0
+ model.save_pretrained(str(output_dir / "best_model"))
+ tokenizer.save_pretrained(str(output_dir / "best_model"))
+ logger.info(" New best! macro_f1=%.4f saved to best_model/", best_f1)
+ if "confusion_matrix" in val_metrics:
+ cm = np.array(val_metrics["confusion_matrix"])
+ plot_confusion_matrix(cm, output_dir / f"cm_epoch{epoch}.png", epoch)
+ else:
+ patience_counter += 1
+
+ # Save last (for resume)
+ model.save_pretrained(str(output_dir / "last_model"))
+ tokenizer.save_pretrained(str(output_dir / "last_model"))
+ with open(output_dir / "last_checkpoint.json", "w") as f:
+ json.dump({
+ "epoch": epoch, "best_f1": best_f1,
+ "patience_counter": patience_counter,
+ "training_log": training_log,
+ }, f, indent=2)
+
+ with open(output_dir / "training_log.json", "w") as f:
+ json.dump(training_log, f, indent=2)
+
+ if patience_counter >= args.patience:
+ logger.info("Early stopping at epoch %d (patience=%d)", epoch, args.patience)
+ break
+
+ if device == "cuda":
+ torch.cuda.empty_cache()
+
+ # Save config
+ config = {
+ "base_model": args.model_id,
+ "method": "PEFT LoRA",
+ "lora_r": args.lora_r,
+ "lora_alpha": args.lora_alpha,
+ "target_modules": ["query", "value"],
+ "num_classes": NUM_CLASSES,
+ "labels": LABELS_7CLASS,
+ "best_val_f1": best_f1,
+ "training_args": vars(args),
+ }
+ with open(output_dir / "config.json", "w") as f:
+ json.dump(config, f, indent=2, ensure_ascii=False)
+
+ logger.info("Training complete. Best F1=%.4f", best_f1)
+
+
+def main():
+ parser = argparse.ArgumentParser(description="LoRA fine-tune KcELECTRA 7-class")
+ parser.add_argument("--train-manifest", required=True)
+ parser.add_argument("--val-manifest", required=True)
+ parser.add_argument("--output-dir", default="data/models/lora_kcelectra_7class")
+ parser.add_argument("--model-id", default="beomi/KcELECTRA-base-v2022")
+ parser.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
+ parser.add_argument("--epochs", type=int, default=10)
+ parser.add_argument("--batch-size", type=int, default=16)
+ parser.add_argument("--accumulate-steps", type=int, default=2)
+ parser.add_argument("--lr", type=float, default=2e-4)
+ parser.add_argument("--weight-decay", type=float, default=0.01)
+ parser.add_argument("--patience", type=int, default=3)
+ parser.add_argument("--max-length", type=int, default=128)
+ parser.add_argument("--lora-r", type=int, default=16)
+ parser.add_argument("--lora-alpha", type=int, default=32)
+ parser.add_argument("--lora-dropout", type=float, default=0.1)
+ parser.add_argument("--resume", action="store_true")
+ args = parser.parse_args()
+
+ torch.manual_seed(42)
+ random.seed(42)
+ np.random.seed(42)
+
+ train(args)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/train_whisper_emotion_head.py b/scripts/train_whisper_emotion_head.py
new file mode 100644
index 0000000000000000000000000000000000000000..09711821316abe737c534925e1dd2f280c6d4d8f
--- /dev/null
+++ b/scripts/train_whisper_emotion_head.py
@@ -0,0 +1,219 @@
+#!/usr/bin/env python3
+"""Whisper-Medium Encoder + Linear Emotion Head ํ์ต.
+
+Whisper encoder๋ฅผ freezeํ๊ณ linear classifier head๋ง ํ์ตํ์ฌ
+benchmark_ser_models.py์ WhisperMediumAdapter์ ์ฌ์ฉํ ์ฒดํฌํฌ์ธํธ๋ฅผ ์์ฑํ๋ค.
+
+Usage:
+ # AI Hub ํ
์คํธ ์๋ธ์
์ธ์ ๋ฐ์ดํฐ๋ก ํ์ต (test leakage ๋ฐฉ์ง)
+ python scripts/train_whisper_emotion_head.py \\
+ --train-dir data/evaluation/korean/train_audio \\
+ --val-dir data/evaluation/korean/val_audio \\
+ --output data/models/whisper_emotion_head.pt
+
+ # prepare_aihub_test_subset.py์ ์ถ๋ ฅ์ผ๋ก quick test (test leakage ์ฃผ์)
+ python scripts/train_whisper_emotion_head.py \\
+ --train-dir data/evaluation/korean/test_audio \\
+ --epochs 3 --output data/models/whisper_emotion_head.pt
+"""
+
+from __future__ import annotations
+
+import argparse
+import glob
+import logging
+import os
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, Dataset
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+EVAL_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear"]
+LABEL_TO_IDX = {label: i for i, label in enumerate(EVAL_LABELS)}
+
+
+class EmotionAudioDataset(Dataset):
+ """Load WAV files organized by emotion class directory."""
+
+ def __init__(self, root_dir: str, processor, max_samples_per_class: int | None = None):
+ self.samples = []
+ self.processor = processor
+
+ for label in EVAL_LABELS:
+ class_dir = Path(root_dir) / label
+ if not class_dir.exists():
+ logger.warning("Class directory not found: %s", class_dir)
+ continue
+
+ wavs = sorted(glob.glob(str(class_dir / "*.wav")))
+ if max_samples_per_class and len(wavs) > max_samples_per_class:
+ wavs = wavs[:max_samples_per_class]
+
+ for wav_path in wavs:
+ self.samples.append({
+ "path": wav_path,
+ "label": LABEL_TO_IDX[label],
+ })
+
+ logger.info("Dataset: %d samples from %s", len(self.samples), root_dir)
+
+ def __len__(self):
+ return len(self.samples)
+
+ def __getitem__(self, idx):
+ sample = self.samples[idx]
+ import librosa
+ audio, sr = librosa.load(sample["path"], sr=16000)
+ inputs = self.processor(audio, sampling_rate=16000, return_tensors="pt")
+ features = inputs.input_features.squeeze(0) # (n_mels, T)
+ return features, sample["label"]
+
+
+def collate_fn(batch):
+ features, labels = zip(*batch)
+ # Pad features to same length
+ max_len = max(f.shape[1] for f in features)
+ padded = []
+ for f in features:
+ if f.shape[1] < max_len:
+ pad = torch.zeros(f.shape[0], max_len - f.shape[1])
+ f = torch.cat([f, pad], dim=1)
+ padded.append(f)
+ return torch.stack(padded), torch.tensor(labels, dtype=torch.long)
+
+
+def train(args):
+ from transformers import WhisperModel, WhisperFeatureExtractor
+
+ device = torch.device(args.device)
+
+ # Load Whisper encoder (frozen)
+ logger.info("Loading Whisper-Medium encoder...")
+ processor = WhisperFeatureExtractor.from_pretrained("openai/whisper-medium")
+ whisper = WhisperModel.from_pretrained("openai/whisper-medium").to(device)
+ whisper.eval()
+ for param in whisper.parameters():
+ param.requires_grad = False
+
+ hidden_dim = whisper.config.d_model # 1024
+ head = nn.Linear(hidden_dim, len(EVAL_LABELS)).to(device)
+
+ # Dataset
+ train_dataset = EmotionAudioDataset(args.train_dir, processor, args.max_samples_per_class)
+ if len(train_dataset) == 0:
+ logger.error("No training samples found")
+ return
+
+ train_loader = DataLoader(
+ train_dataset, batch_size=args.batch_size, shuffle=True,
+ collate_fn=collate_fn, num_workers=0,
+ )
+
+ val_loader = None
+ if args.val_dir and Path(args.val_dir).exists():
+ val_dataset = EmotionAudioDataset(args.val_dir, processor)
+ if len(val_dataset) > 0:
+ val_loader = DataLoader(
+ val_dataset, batch_size=args.batch_size, shuffle=False,
+ collate_fn=collate_fn, num_workers=0,
+ )
+
+ # Optimizer
+ optimizer = torch.optim.Adam(head.parameters(), lr=args.lr)
+ criterion = nn.CrossEntropyLoss()
+
+ # Training loop
+ best_val_acc = 0.0
+ for epoch in range(args.epochs):
+ head.train()
+ total_loss = 0
+ correct = 0
+ total = 0
+
+ for batch_idx, (features, labels) in enumerate(train_loader):
+ features = features.to(device)
+ labels = labels.to(device)
+
+ with torch.no_grad():
+ encoder_out = whisper.encoder(features)
+ hidden = encoder_out.last_hidden_state # (B, T, D)
+ pooled = hidden.mean(dim=1) # (B, D)
+
+ logits = head(pooled)
+ loss = criterion(logits, labels)
+
+ optimizer.zero_grad()
+ loss.backward()
+ optimizer.step()
+
+ total_loss += loss.item() * labels.size(0)
+ preds = logits.argmax(dim=1)
+ correct += (preds == labels).sum().item()
+ total += labels.size(0)
+
+ train_acc = correct / max(total, 1)
+ avg_loss = total_loss / max(total, 1)
+ logger.info("Epoch %d/%d: loss=%.4f, train_acc=%.3f",
+ epoch + 1, args.epochs, avg_loss, train_acc)
+
+ # Validation
+ if val_loader:
+ head.eval()
+ val_correct = 0
+ val_total = 0
+ with torch.no_grad():
+ for features, labels in val_loader:
+ features = features.to(device)
+ labels = labels.to(device)
+ encoder_out = whisper.encoder(features)
+ pooled = encoder_out.last_hidden_state.mean(dim=1)
+ logits = head(pooled)
+ preds = logits.argmax(dim=1)
+ val_correct += (preds == labels).sum().item()
+ val_total += labels.size(0)
+ val_acc = val_correct / max(val_total, 1)
+ logger.info(" val_acc=%.3f", val_acc)
+
+ if val_acc > best_val_acc:
+ best_val_acc = val_acc
+ save_checkpoint(head, args.output, epoch, val_acc)
+ else:
+ # No val set โ save latest
+ save_checkpoint(head, args.output, epoch, train_acc)
+
+ logger.info("Training complete. Best checkpoint: %s", args.output)
+
+
+def save_checkpoint(head: nn.Linear, path: str, epoch: int, accuracy: float):
+ Path(path).parent.mkdir(parents=True, exist_ok=True)
+ torch.save(head.state_dict(), path)
+ logger.info("Saved checkpoint: %s (epoch=%d, acc=%.3f)", path, epoch + 1, accuracy)
+
+
+def main():
+ parser = argparse.ArgumentParser(description="Train Whisper emotion classifier head")
+ parser.add_argument("--train-dir", required=True,
+ help="ํ์ต ์ค๋์ค ๋๋ ํ ๋ฆฌ ({emotion}/*.wav ๊ตฌ์กฐ)")
+ parser.add_argument("--val-dir", default=None,
+ help="๊ฒ์ฆ ์ค๋์ค ๋๋ ํ ๋ฆฌ (์์ผ๋ฉด train accuracy๋ก ํ๋จ)")
+ parser.add_argument("--output", default="data/models/whisper_emotion_head.pt",
+ help="์ถ๋ ฅ ์ฒดํฌํฌ์ธํธ ๊ฒฝ๋ก")
+ parser.add_argument("--epochs", type=int, default=10)
+ parser.add_argument("--batch-size", type=int, default=8)
+ parser.add_argument("--lr", type=float, default=1e-3)
+ parser.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
+ parser.add_argument("--max-samples-per-class", type=int, default=None,
+ help="ํด๋์ค๋น ์ต๋ ํ์ต ์ํ ์")
+ args = parser.parse_args()
+
+ train(args)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/validate_text_emotion_english.py b/scripts/validate_text_emotion_english.py
new file mode 100644
index 0000000000000000000000000000000000000000..3b9e06dfe3d0de1f6be2e61873d22af780465f34
--- /dev/null
+++ b/scripts/validate_text_emotion_english.py
@@ -0,0 +1,137 @@
+#!/usr/bin/env python3
+"""์์ด ํ
์คํธ ๊ฐ์ ๋ชจ๋ธ (DistilRoBERTa) sanity check.
+
+hand-crafted ์์ด ๋ฌธ์ฅ์ผ๋ก ๋ชจ๋ธ ๋ก๋ฉ, ์ถ๋ ฅ ํฌ๋งท, ๊ธฐ๋ณธ ์ ํ๋๋ฅผ ๊ฒ์ฆํ๋ค.
+
+Usage:
+ python scripts/validate_text_emotion_english.py
+"""
+
+from __future__ import annotations
+
+import logging
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).parent.parent
+sys.path.insert(0, str(PROJECT_ROOT))
+
+logging.basicConfig(
+ level=logging.INFO,
+ format="%(asctime)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger("validate_text_en")
+
+PROJECT_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear", "disgust"]
+
+# (text, expected_emotion) โ ๋ช
ํํ ๊ฐ์ ํํ ๋ฌธ์ฅ
+TEST_CASES = [
+ # joy
+ ("I'm so happy today, everything is wonderful!", "joy"),
+ ("This is the best day of my life!", "joy"),
+ ("I'm thrilled about the good news!", "joy"),
+ # anger
+ ("This makes me absolutely furious!", "anger"),
+ ("I can't believe how unfair this is, I'm so angry!", "anger"),
+ ("Stop doing that, it's driving me crazy!", "anger"),
+ # sadness
+ ("I feel so sad and lonely right now.", "sadness"),
+ ("I miss you so much, it hurts.", "sadness"),
+ ("I can't stop crying, everything feels hopeless.", "sadness"),
+ # fear
+ ("I'm really scared, something is wrong.", "fear"),
+ ("I'm terrified of what might happen next.", "fear"),
+ ("Help me, I'm so afraid!", "fear"),
+ # surprise
+ ("Oh my god, I can't believe it!", "surprise"),
+ ("Wow, I never expected that to happen!", "surprise"),
+ ("What?! That's absolutely incredible!", "surprise"),
+ # neutral
+ ("The meeting is scheduled for 3pm.", "neutral"),
+ ("I need to buy groceries on the way home.", "neutral"),
+ ("The temperature today is around 20 degrees.", "neutral"),
+ # disgust
+ ("That's absolutely disgusting, I feel sick.", "disgust"),
+ ("This food tastes terrible, it's revolting.", "disgust"),
+ ("I can't stand the smell, it's nauseating.", "disgust"),
+]
+
+
+def validate():
+ """DistilRoBERTa ๋ชจ๋ธ ๊ฒ์ฆ."""
+ from src.stage2.text_emotion import predict as text_predict
+
+ logger.info("=" * 60)
+ logger.info("์์ด ํ
์คํธ ๊ฐ์ ๋ชจ๋ธ (DistilRoBERTa) ๊ฒ์ฆ ์์")
+ logger.info("=" * 60)
+
+ passed = 0
+ failed = 0
+ results = []
+
+ for text, expected in TEST_CASES:
+ result = text_predict(text, language="en")
+
+ # ์ถ๋ ฅ ํฌ๋งท ๊ฒ์ฆ
+ assert "emotion" in result, f"Missing 'emotion' key in result"
+ assert "confidence" in result, f"Missing 'confidence' key in result"
+ assert "scores" in result, f"Missing 'scores' key in result"
+
+ # ๋ชจ๋ ํ๋ก์ ํธ ๋ผ๋ฒจ์ด scores์ ์๋์ง ํ์ธ
+ for label in PROJECT_LABELS:
+ assert label in result["scores"], f"Missing label '{label}' in scores"
+
+ # scores ํฉ๊ณ ~1.0 ํ์ธ
+ score_sum = sum(result["scores"].values())
+ assert abs(score_sum - 1.0) < 0.01, f"Scores sum={score_sum}, expected ~1.0"
+
+ predicted = result["emotion"]
+ confidence = result["confidence"]
+ match = predicted == expected
+
+ if match:
+ passed += 1
+ status = "PASS"
+ else:
+ failed += 1
+ status = "FAIL"
+
+ results.append({
+ "text": text[:50],
+ "expected": expected,
+ "predicted": predicted,
+ "confidence": confidence,
+ "match": match,
+ })
+
+ logger.info(
+ f" [{status}] expected={expected:10s} predicted={predicted:10s} "
+ f"conf={confidence:.3f} | \"{text[:45]}...\""
+ if len(text) > 45 else
+ f" [{status}] expected={expected:10s} predicted={predicted:10s} "
+ f"conf={confidence:.3f} | \"{text}\""
+ )
+
+ total = passed + failed
+ accuracy = passed / total if total > 0 else 0
+
+ logger.info(f"\n{'='*60}")
+ logger.info(f"๊ฒฐ๊ณผ: {passed}/{total} PASS ({accuracy:.1%})")
+ logger.info(f" Passed: {passed}")
+ logger.info(f" Failed: {failed}")
+
+ if accuracy >= 0.8:
+ logger.info("ํ์ : PASS โ ๋ชจ๋ธ์ด ์ ์์ ์ผ๋ก ๋์ํฉ๋๋ค.")
+ elif accuracy >= 0.6:
+ logger.info("ํ์ : WARN โ ์ผ๋ถ ๊ฐ์ ์์ ๋ถ์ ํํฉ๋๋ค.")
+ else:
+ logger.info("ํ์ : FAIL โ ๋ชจ๋ธ์ ๋ฌธ์ ๊ฐ ์์ ์ ์์ต๋๋ค.")
+
+ logger.info(f"{'='*60}")
+
+ return accuracy >= 0.6
+
+
+if __name__ == "__main__":
+ success = validate()
+ sys.exit(0 if success else 1)
diff --git a/src/common/__init__.py b/src/common/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/src/common/constants.py b/src/common/constants.py
new file mode 100644
index 0000000000000000000000000000000000000000..a7c64cca0a86fb393346fc9480ef991c1a3c7b15
--- /dev/null
+++ b/src/common/constants.py
@@ -0,0 +1,59 @@
+# ๊ณต์ฉ ์์
+
+EMOTION_LABELS = [
+ "neutral",
+ "joy",
+ "sadness",
+ "anger",
+ "surprise",
+ "fear",
+ "disgust",
+]
+
+EMOTION_TO_VALENCE = {
+ "joy": 1.0,
+ "surprise": 0.5,
+ "neutral": 0.0,
+ "sadness": -0.5,
+ "fear": -0.6,
+ "anger": -0.8,
+ "disgust": -0.9,
+}
+
+FUSION_WEIGHTS = {
+ "audio": 0.6,
+ "text": 0.4,
+}
+
+# Korean emotion-specific fusion weights โ trained via gradient descent
+# Parameterization: w_a = sigmoid(ฮฑ), w_t = 1 - w_a (7 params jointly optimized)
+# Data: AI Hub 263 val (1,294 samples, stratified 80/20 train/val)
+# Val macro F1 = 0.8724 (vs 0.8744 fixed 60/40, 0.8757 greedy v1).
+# Full 263 val macro F1 = 0.8748 (vs 0.8736 greedy v1, 0.8347 fixed 60/40).
+# Text-dominant pattern preserved (Korean KcELECTRA LoRA > audio LoRA on most classes).
+EMOTION_FUSION_WEIGHTS_KO = {
+ "neutral": {"audio": 0.53, "text": 0.47},
+ "joy": {"audio": 0.25, "text": 0.75},
+ "sadness": {"audio": 0.16, "text": 0.84},
+ "anger": {"audio": 0.11, "text": 0.89},
+ "surprise": {"audio": 0.22, "text": 0.78},
+ "fear": {"audio": 0.10, "text": 0.90},
+ "disgust": {"audio": 0.15, "text": 0.85},
+}
+
+# English emotion-specific fusion weights โ trained via gradient descent
+# Parameterization: w_a = sigmoid(ฮฑ), w_t = 1 - w_a (7 params jointly optimized)
+# Data: JL-Corpus + SAVEE + MELD + RAVDESS phone (fear/disgust/sadness) = 2,821 samples
+# Val macro F1 = 0.7596 (vs 0.7473 fixed 60/40, 0.7114 audio-only, 0.6295 greedy v1)
+EMOTION_FUSION_WEIGHTS_EN = {
+ "neutral": {"audio": 0.78, "text": 0.22},
+ "joy": {"audio": 0.61, "text": 0.39},
+ "sadness": {"audio": 0.70, "text": 0.30},
+ "anger": {"audio": 0.60, "text": 0.40},
+ "surprise": {"audio": 0.46, "text": 0.54},
+ "fear": {"audio": 0.74, "text": 0.26},
+ "disgust": {"audio": 0.77, "text": 0.23},
+}
+
+# Default (Korean) โ back-compat alias
+EMOTION_FUSION_WEIGHTS = EMOTION_FUSION_WEIGHTS_KO
diff --git a/src/common/phone_simulator.py b/src/common/phone_simulator.py
new file mode 100644
index 0000000000000000000000000000000000000000..21aa702389cb68cc3e37f44e3425b8c5353a2ffe
--- /dev/null
+++ b/src/common/phone_simulator.py
@@ -0,0 +1,130 @@
+"""
+์ ํ ํตํ ํ์ง ์๋ฎฌ๋ ์ด์
(PSTN)
+
+๊นจ๋ํ ๋
น์ ์ค๋์ค๋ฅผ ์ ํ ํตํ ํ์ง๋ก ๋ณํํ์ฌ
+AI Hub ๋ฑ์ ์คํ๋์ค ๋
น์ ๋ฐ์ดํฐ๋ฅผ ํ์ต์ฉ ํตํ ๋ฐ์ดํฐ๋ก ์ ์ฒ๋ฆฌํ๋ค.
+
+3๋จ๊ณ ์ฒ๋ฆฌ:
+ 1. ๋ฐด๋ํจ์ค ํํฐ๋ง (300โ3400 Hz) โ ITU-T G.712
+ 2. 8 kHz ๋ค์ด์ํ๋ง (anti-aliasing ํฌํจ)
+ 3. G.711 ๋น์ ํ ์์ํ (A-law / ฮผ-law companding)
+"""
+
+import audioop
+import random
+from enum import Enum
+
+import numpy as np
+import scipy.signal as signal
+
+
+class CompandingType(str, Enum):
+ ALAW = "alaw" # ํ๊ตญ/์ ๋ฝ/์์์ PSTN ํ์ค
+ ULAW = "ulaw" # ๋ถ๋ฏธ PSTN ํ์ค
+ RANDOM = "random" # ๋๋ค ์ ํ (ํ์ต ๋ฐ์ดํฐ ๋ค์์ฑ ํ๋ณด)
+
+
+class PhoneSimulator:
+ """๊นจ๋ํ ๋
น์ ์ค๋์ค โ ์ ํ ํตํ ํ์ง ๋ณํ๊ธฐ"""
+
+ # PSTN ํ์ค ํ๋ผ๋ฏธํฐ
+ PSTN_LOW_FREQ = 300.0 # Hz โ ITU-T G.712 ํํ
+ PSTN_HIGH_FREQ = 3400.0 # Hz โ ITU-T G.712 ์ํ
+ PSTN_SAMPLE_RATE = 8000 # Hz โ G.711 ํ์ค ์ํ๋ ์ดํธ
+ FILTER_ORDER = 5 # Butterworth ํํฐ ์ฐจ์
+
+ def __init__(self, companding: CompandingType = CompandingType.RANDOM):
+ """
+ Args:
+ companding: ์์ํ ๋ฐฉ์. RANDOM์ด๋ฉด ํ์ผ๋ง๋ค alaw/ulaw ๋๋ค ์ ํ
+ """
+ self.companding = companding
+
+ def process(self, audio: np.ndarray, sr: int) -> tuple[np.ndarray, int]:
+ """
+ ์ ํ ํตํ ํ์ง ์๋ฎฌ๋ ์ด์
์ ์ฉ.
+
+ Args:
+ audio: float32 mono numpy array (๋ฒ์: -1.0 ~ 1.0)
+ sr: ์๋ณธ ์ํ๋ ์ดํธ
+
+ Returns:
+ (์ฒ๋ฆฌ๋ ์ค๋์ค, ์ ์ํ๋ ์ดํธ=8000)
+ """
+ if audio.ndim != 1:
+ raise ValueError(f"Mono audio expected, got shape {audio.shape}")
+
+ # Step 1: ๋ฐด๋ํจ์ค ํํฐ๋ง (300โ3400 Hz)
+ audio = self._bandpass_filter(audio, sr)
+
+ # Step 2: 8 kHz ๋ค์ด์ํ๋ง
+ audio = self._downsample(audio, sr)
+
+ # Step 3: G.711 companding (encodeโdecode round-trip)
+ audio = self._compand(audio)
+
+ return audio, self.PSTN_SAMPLE_RATE
+
+ def _bandpass_filter(self, audio: np.ndarray, sr: int) -> np.ndarray:
+ """
+ ITU-T G.712 ๋์ญ ํํฐ๋ง.
+
+ 300 Hz ๋ฏธ๋ง(ํ ๋
ธ์ด์ฆ) + 3400 Hz ์ด์(์น์ฐฐ์) ์ ๊ฑฐ.
+ 5์ฐจ Butterworth: ์ถฉ๋ถํ ๊ฐํ๋ฅด๋ฉด์ ringing ์ต์ํ.
+ """
+ nyq = sr / 2.0
+ low = self.PSTN_LOW_FREQ / nyq
+ high = self.PSTN_HIGH_FREQ / nyq
+
+ # ๋์ดํด์คํธ ์ด์์ด๋ฉด ํํฐ ์ ์ฉ ๋ถ๊ฐ (์ด๋ฏธ ๋์ญ ๋ด)
+ if high >= 1.0:
+ high = 0.99
+ if low <= 0.0:
+ low = 0.01
+
+ b, a = signal.butter(self.FILTER_ORDER, [low, high], btype="band")
+ return signal.filtfilt(b, a, audio).astype(np.float32)
+
+ def _downsample(self, audio: np.ndarray, sr: int) -> np.ndarray:
+ """
+ Anti-aliasing + ๋ค์ด์ํ๋ง.
+
+ scipy.signal.resample_poly๋ ๋ด๋ถ์ ์ผ๋ก anti-aliasing ํํฐ๋ฅผ ์ ์ฉํ์ฌ
+ ์์ผ๋ฆฌ์ด์ฑ์ ๋ฐฉ์งํ๋ค.
+ """
+ if sr == self.PSTN_SAMPLE_RATE:
+ return audio
+
+ # GCD ๊ธฐ๋ฐ rational resampling (resample_poly๊ฐ ๋ ์ ํ)
+ gcd = np.gcd(sr, self.PSTN_SAMPLE_RATE)
+ up = self.PSTN_SAMPLE_RATE // gcd
+ down = sr // gcd
+ return signal.resample_poly(audio, up, down).astype(np.float32)
+
+ def _compand(self, audio: np.ndarray) -> np.ndarray:
+ """
+ G.711 A-law/ฮผ-law encodeโdecode round-trip.
+
+ 16๋นํธ โ 8๋นํธ ์์ถ โ 16๋นํธ ๋ณต์ ๊ณผ์ ์์
+ ๋น์ ํ ์์ํ ๋
ธ์ด์ฆ๊ฐ ๋ฐ์ํ์ฌ ์ ํ ํน์ ์ '๊ฑฐ์น' ์์์ ๋ง๋ ๋ค.
+ """
+ # companding ๋ฐฉ์ ๊ฒฐ์
+ if self.companding == CompandingType.RANDOM:
+ method = random.choice([CompandingType.ALAW, CompandingType.ULAW])
+ else:
+ method = self.companding
+
+ # float32 โ 16-bit PCM
+ pcm16 = np.clip(audio * 32767, -32768, 32767).astype(np.int16)
+ raw_bytes = pcm16.tobytes()
+
+ # encode (16bit โ 8bit) โ decode (8bit โ 16bit) round-trip
+ if method == CompandingType.ALAW:
+ compressed = audioop.lin2alaw(raw_bytes, 2)
+ decompressed = audioop.alaw2lin(compressed, 2)
+ else:
+ compressed = audioop.lin2ulaw(raw_bytes, 2)
+ decompressed = audioop.ulaw2lin(compressed, 2)
+
+ # 16-bit PCM โ float32
+ return np.frombuffer(decompressed, dtype=np.int16).astype(np.float32) / 32767.0
diff --git a/src/common/schemas.py b/src/common/schemas.py
new file mode 100644
index 0000000000000000000000000000000000000000..b6d9ff23579e464f2157f679d08609b375be0570
--- /dev/null
+++ b/src/common/schemas.py
@@ -0,0 +1,97 @@
+# Stage ๊ฐ JSON ์ธํฐํ์ด์ค (๊ณ์ฝ)
+# ๋ณ๊ฒฝ ์ 3์ธ ํฉ์ ํ์
+
+from pydantic import BaseModel, Field
+
+
+# --- Stage 1 Models ---
+
+class Models(BaseModel):
+ diarization: str
+ asr: str
+ language_id: str
+ alignment: str
+
+
+class ProcessingInfo(BaseModel):
+ processing_time_sec: float = Field(gt=0)
+ models: Models
+ device: str # "cpu" | "cuda" | "cuda:N"
+ language_chosen: str | None = None # "ko" | "en", text-based decision used for alignment
+ korean_ratio: float | None = Field(default=None, ge=0.0, le=1.0)
+
+
+class Segment(BaseModel):
+ segment_id: int = Field(ge=0)
+ speaker_id: str
+ start: float = Field(ge=0.0)
+ end: float = Field(gt=0.0)
+ text: str
+ language: str # "ko" | "en"
+ audio_path: str
+ confidence: float = Field(ge=0.0, le=1.0)
+
+
+class Stage1Output(BaseModel):
+ call_id: str
+ duration: float = Field(gt=0.0)
+ speakers: list[str] # ["speaker_0", "speaker_1"]
+ audio_path: str # path to original audio file
+ segments: list[Segment]
+ processing_info: ProcessingInfo
+
+
+class EmotionResult(BaseModel):
+ speaker_id: str
+ segment_id: int
+ audio_emotion: str
+ audio_confidence: float
+ text_emotion: str
+ text_confidence: float
+ fused_emotion: str
+ fused_confidence: float
+
+
+class SpeakerSummary(BaseModel):
+ dominant_emotion: str
+ emotion_distribution: dict[str, float]
+ avg_confidence: float
+
+
+class Stage2Output(BaseModel):
+ call_id: str
+ emotions: list[EmotionResult]
+ speaker_summaries: dict[str, SpeakerSummary]
+
+
+class CharacterReaction(BaseModel):
+ speaker_id: str
+ solo_state: str
+ pair_state: str
+ role: str | None = None # "giver" | "receiver" | None
+ # Comforting strength (1..3). Default 2 preserves the P0 baseline (3 cycles
+ # to heal). Higher = stronger comfort โ receiver heals in fewer cycles.
+ # Client mapping lives in app/src/hooks/healingThreshold.ts.
+ intensity: int = Field(default=2, ge=1, le=3)
+
+
+class GardenUpdate(BaseModel):
+ growth_delta: int
+ total_level: int
+ mood: str
+
+
+class RecapCard(BaseModel):
+ title: str
+ summary: str
+ highlights: list[str]
+
+
+class Stage3Output(BaseModel):
+ call_id: str
+ character_reactions: list[CharacterReaction]
+ garden_update: GardenUpdate
+ recap_card: RecapCard
+ # Stage2 data surfaced at top level for frontend convenience
+ emotions: list[EmotionResult] = Field(default_factory=list)
+ speaker_summaries: dict[str, SpeakerSummary] = Field(default_factory=dict)
diff --git a/src/stage1/__init__.py b/src/stage1/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..15cc3dc5ae78b0e5f26e6c5e7a3f9d6880022dd0
--- /dev/null
+++ b/src/stage1/__init__.py
@@ -0,0 +1,5 @@
+"""Stage 1: Speaker Diarization + ASR."""
+
+from src.stage1.process import process
+
+__all__ = ["process"]
diff --git a/src/stage1/diarization.py b/src/stage1/diarization.py
new file mode 100644
index 0000000000000000000000000000000000000000..074c6c35411e35a7d3067442864ac22c62a22d96
--- /dev/null
+++ b/src/stage1/diarization.py
@@ -0,0 +1,125 @@
+"""Stage 1 โ Speaker Diarization Module.
+
+Uses pyannote-audio 4.x (speaker-diarization-3.1) to identify who spoke when.
+"""
+
+import logging
+import os
+from dataclasses import dataclass
+
+import torch
+import torchaudio
+from pyannote.audio import Pipeline
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class DiarSegment:
+ speaker_id: str
+ start: float
+ end: float
+
+
+_pipeline: Pipeline | None = None
+
+
+def _load_pipeline(model: str, device: str) -> Pipeline:
+ """Load and cache the pyannote diarization pipeline."""
+ global _pipeline
+ if _pipeline is not None:
+ return _pipeline
+
+ hf_token = os.environ.get("HF_TOKEN")
+ if not hf_token:
+ raise RuntimeError(
+ "HF_TOKEN environment variable required for pyannote model download. "
+ "Set it with: export HF_TOKEN=your_huggingface_token"
+ )
+
+ logger.info("Loading pyannote pipeline: %s", model)
+ _pipeline = Pipeline.from_pretrained(model, token=hf_token)
+
+ torch_device = torch.device(device if torch.cuda.is_available() and "cuda" in device else "cpu")
+ _pipeline.to(torch_device)
+ logger.info("Pyannote pipeline loaded on %s", torch_device)
+
+ return _pipeline
+
+
+def _merge_adjacent(
+ segments: list[DiarSegment], gap_threshold: float
+) -> list[DiarSegment]:
+ """Merge same-speaker segments closer than gap_threshold seconds."""
+ if not segments:
+ return segments
+
+ merged: list[DiarSegment] = [segments[0]]
+ for seg in segments[1:]:
+ prev = merged[-1]
+ if seg.speaker_id == prev.speaker_id and (seg.start - prev.end) < gap_threshold:
+ prev.end = max(prev.end, seg.end)
+ else:
+ merged.append(seg)
+ return merged
+
+
+def diarize(audio_path: str, config: dict) -> list[DiarSegment]:
+ """Run speaker diarization on audio.
+
+ Args:
+ audio_path: Path to 16kHz mono WAV file.
+ config: stage1.diarization config section.
+
+ Returns:
+ List of DiarSegment sorted by start time.
+ """
+ model = config.get("model", "pyannote/speaker-diarization-3.1")
+ num_speakers = config.get("num_speakers", 2)
+ merge_gap = config.get("merge_gap_sec", 0.15)
+ device = config.get("device", "cuda:0")
+
+ pipeline = _load_pipeline(model, device)
+
+ logger.info("Running diarization (num_speakers=%d) on %s", num_speakers, audio_path)
+
+ # pyannote 4.x: pass in-memory waveform to avoid torchcodec dependency
+ waveform, sample_rate = torchaudio.load(audio_path)
+ audio_dict = {"waveform": waveform, "sample_rate": sample_rate}
+ output = pipeline(audio_dict, num_speakers=num_speakers)
+
+ # pyannote 4.x returns DiarizeOutput with .speaker_diarization Annotation
+ if hasattr(output, 'speaker_diarization'):
+ annotation = output.speaker_diarization
+ elif hasattr(output, 'itertracks'):
+ annotation = output
+ else:
+ raise RuntimeError(f"Unexpected pyannote output type: {type(output)}")
+
+ segments: list[DiarSegment] = []
+ for turn, _, speaker in annotation.itertracks(yield_label=True):
+ segments.append(DiarSegment(
+ speaker_id=speaker,
+ start=round(turn.start, 3),
+ end=round(turn.end, 3),
+ ))
+
+ segments.sort(key=lambda s: s.start)
+ segments = _merge_adjacent(segments, merge_gap)
+
+ # Normalize speaker labels to speaker_0 / speaker_1
+ speaker_map: dict[str, str] = {}
+ for seg in segments:
+ if seg.speaker_id not in speaker_map:
+ speaker_map[seg.speaker_id] = f"speaker_{len(speaker_map)}"
+ seg.speaker_id = speaker_map[seg.speaker_id]
+
+ logger.info(
+ "Diarization complete: %d segments, %d speakers",
+ len(segments), len(speaker_map),
+ )
+
+ if len(speaker_map) < 2:
+ logger.warning("Only %d speaker(s) detected", len(speaker_map))
+
+ return segments
\ No newline at end of file
diff --git a/src/stage1/language_id.py b/src/stage1/language_id.py
new file mode 100644
index 0000000000000000000000000000000000000000..c2197bda14ef4d41e26fcd50baab00cbaa019e14
--- /dev/null
+++ b/src/stage1/language_id.py
@@ -0,0 +1,90 @@
+"""Stage 1 โ Language Identification Module.
+
+Decides language from transcribed text using a Korean-character ratio.
+The same 50% majority rule is applied at both the global level
+(by the orchestrator, to pick the wav2vec2 alignment model) and the
+per-segment level (here, to route Stage 2 text emotion).
+
+SenseVoice is reserved for potential future use (emotion2vec fallback).
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+
+logger = logging.getLogger(__name__)
+
+_KOREAN_PATTERN = re.compile(r"[๊ฐ-ํฏ]")
+
+KOREAN_THRESHOLD = 0.5
+MIN_CHARS_FOR_TEXT_DECISION = 10
+
+
+def compute_korean_ratio(text: str) -> tuple[float, int]:
+ """Return (korean_char_ratio, non_space_char_count).
+
+ Empty or whitespace-only text โ (0.0, 0).
+ """
+ stripped = text.replace(" ", "")
+ total = len(stripped)
+ if total == 0:
+ return 0.0, 0
+ korean = len(_KOREAN_PATTERN.findall(text))
+ return korean / total, total
+
+
+def decide_language_from_text(
+ text: str,
+ fallback: str = "ko",
+ threshold: float = KOREAN_THRESHOLD,
+ min_chars: int = MIN_CHARS_FOR_TEXT_DECISION,
+) -> tuple[str, float]:
+ """Return (language, korean_ratio) using the majority rule.
+
+ Falls back to ``fallback`` when the transcript has fewer than
+ ``min_chars`` non-space characters โ short text isn't statistically
+ meaningful and a single character can swing the ratio wildly.
+ """
+ ratio, total = compute_korean_ratio(text)
+ if total < min_chars:
+ return fallback, ratio
+ return ("ko" if ratio >= threshold else "en"), ratio
+
+
+def detect_languages(
+ segments: list[dict],
+ audio_path: str,
+ config: dict,
+ device: str,
+ whisper_language: str = "ko",
+) -> list[str]:
+ """Detect language for each segment using the text-based majority rule.
+
+ Args:
+ segments: List of segment dicts with 'start', 'end', 'text' keys.
+ audio_path: Path to the full audio file (unused, kept for interface).
+ config: stage1.language_id config section.
+ device: Compute device (unused, kept for interface).
+ whisper_language: Whisper's globally detected language, used as
+ fallback when a segment's text is empty or too short.
+
+ Returns:
+ List of language codes ("ko" or "en"), one per segment.
+ """
+ enabled = config.get("enabled", True)
+ if not enabled:
+ logger.info("Language ID disabled, using whisper language: %s", whisper_language)
+ return [whisper_language] * len(segments)
+
+ languages: list[str] = []
+ for seg in segments:
+ text = seg.get("text", "")
+ lang, _ = decide_language_from_text(text, fallback=whisper_language)
+ languages.append(lang)
+
+ logger.info(
+ "Language detection complete: %d segments (whisper=%s)",
+ len(languages), whisper_language,
+ )
+ return languages
diff --git a/src/stage1/process.py b/src/stage1/process.py
new file mode 100644
index 0000000000000000000000000000000000000000..e2b7c78e22d37e5a8a9e26a5bbbeaf1af2de7969
--- /dev/null
+++ b/src/stage1/process.py
@@ -0,0 +1,329 @@
+"""Stage 1 โ Orchestrator.
+
+Entry point: process(audio_path) -> Stage1Output
+
+Pipeline:
+ 1. Audio validation & preprocessing
+ 2. Speaker diarization (pyannote)
+ 3. ASR transcription (Faster-Whisper via WhisperX)
+ 4. Forced alignment (WhisperX wav2vec2)
+ 5. Diarization-transcript merge
+ 6. Language detection (SenseVoice-Small)
+ 7. Segment audio extraction
+ 8. Output assembly
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import time
+import uuid
+from pathlib import Path
+
+import torch
+import torchaudio
+import yaml
+
+from src.common.schemas import (
+ Models,
+ ProcessingInfo,
+ Segment,
+ Stage1Output,
+)
+from src.stage1.diarization import diarize
+from src.stage1.language_id import decide_language_from_text, detect_languages
+from src.stage1.transcription import align, assign_speakers, transcribe
+
+logger = logging.getLogger(__name__)
+
+SUPPORTED_EXTENSIONS = {".wav", ".mp3", ".m4a", ".ogg"}
+
+
+def _load_config(config: dict | None = None) -> dict:
+ """Load config from yaml if not provided."""
+ if config is not None:
+ return config
+
+ config_path = Path(__file__).parent.parent.parent / "config.yaml"
+ if config_path.exists():
+ with open(config_path) as f:
+ return yaml.safe_load(f).get("stage1", {})
+
+ logger.warning("config.yaml not found, using defaults")
+ return {}
+
+
+def _get_device() -> str:
+ """Determine compute device."""
+ return "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def _validate_audio(audio_path: str, config: dict) -> None:
+ """Step 1a: Validate audio file."""
+ path = Path(audio_path)
+
+ if not path.exists():
+ raise FileNotFoundError(f"Audio file not found: {audio_path}")
+
+ if path.suffix.lower() not in SUPPORTED_EXTENSIONS:
+ raise ValueError(
+ f"Unsupported audio format: {path.suffix}. "
+ f"Supported: {SUPPORTED_EXTENSIONS}"
+ )
+
+
+def _preprocess_audio(audio_path: str, config: dict) -> tuple[str, float]:
+ """Step 1b: Convert to 16kHz mono WAV, return (preprocessed_path, duration)."""
+ preprocess_cfg = config.get("preprocessing", {})
+ target_sr = preprocess_cfg.get("target_sample_rate", 16000)
+ min_dur = preprocess_cfg.get("min_duration_sec", 3)
+ max_dur = preprocess_cfg.get("max_duration_sec", 300)
+
+ waveform, sr = torchaudio.load(audio_path)
+ duration = waveform.shape[1] / sr
+
+ if duration < min_dur:
+ raise ValueError(f"Audio too short: {duration:.1f}s (min: {min_dur}s)")
+ if duration > max_dur:
+ raise ValueError(f"Audio too long: {duration:.1f}s (max: {max_dur}s)")
+
+ # Resample if needed
+ if sr != target_sr:
+ resampler = torchaudio.transforms.Resample(sr, target_sr)
+ waveform = resampler(waveform)
+
+ # Stereo to mono
+ if waveform.shape[0] > 1:
+ waveform = waveform.mean(dim=0, keepdim=True)
+
+ # Peak normalization โ handle volume differences across devices
+ peak = waveform.abs().max()
+ if peak > 0:
+ target_peak = preprocess_cfg.get("target_peak", 0.95)
+ waveform = waveform * (target_peak / peak)
+
+ # Save preprocessed WAV
+ preprocessed_path = audio_path
+ if sr != target_sr or Path(audio_path).suffix.lower() != ".wav":
+ preprocessed_dir = Path(config.get("segments_dir", "data/segments"))
+ preprocessed_dir.mkdir(parents=True, exist_ok=True)
+ preprocessed_path = str(preprocessed_dir / f"_preprocessed_{uuid.uuid4().hex[:8]}.wav")
+ torchaudio.save(preprocessed_path, waveform, target_sr)
+
+ return preprocessed_path, duration
+
+
+def _extract_segment_audio(
+ waveform: torch.Tensor,
+ sr: int,
+ segments: list[dict],
+ call_id: str,
+ segments_dir: str,
+) -> list[str]:
+ """Step 7: Extract audio clips for each segment."""
+ seg_dir = Path(segments_dir) / call_id
+ seg_dir.mkdir(parents=True, exist_ok=True)
+
+ paths: list[str] = []
+ for seg in segments:
+ idx = seg["segment_id"]
+ speaker = seg["speaker_id"]
+ start_sample = int(seg["start"] * sr)
+ end_sample = int(seg["end"] * sr)
+
+ # Clamp to valid range
+ start_sample = max(0, start_sample)
+ end_sample = min(waveform.shape[1], end_sample)
+
+ if end_sample <= start_sample:
+ paths.append("")
+ continue
+
+ clip = waveform[:, start_sample:end_sample]
+ filename = f"seg_{idx:03d}_{speaker}.wav"
+ clip_path = str(seg_dir / filename)
+ torchaudio.save(clip_path, clip, sr)
+ paths.append(clip_path)
+
+ return paths
+
+
+def process(audio_path: str, config: dict | None = None) -> Stage1Output:
+ """Stage 1 entry point.
+
+ Args:
+ audio_path: Path to input audio file (.wav/.mp3/.m4a/.ogg)
+ config: Optional config dict. If None, loads from config.yaml.
+
+ Returns:
+ Stage1Output with call_id, duration, and speaker-segmented transcript.
+
+ Raises:
+ FileNotFoundError: Audio file does not exist.
+ ValueError: Unsupported format, too short, or too long.
+ RuntimeError: Audio decode or model inference failure.
+ """
+ start_time = time.time()
+ cfg = _load_config(config)
+ device = _get_device()
+
+ # --- Step 1: Validate & Preprocess ---
+ _validate_audio(audio_path, cfg)
+ preprocessed_path, duration = _preprocess_audio(audio_path, cfg)
+ call_id = f"call_{time.strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:4]}"
+
+ logger.info("Processing %s (%.1fs) as %s", audio_path, duration, call_id)
+
+ # Load waveform for segment extraction later
+ waveform, sr = torchaudio.load(preprocessed_path)
+
+ # --- Step 2: Speaker Diarization ---
+ diar_cfg = cfg.get("diarization", {})
+ diar_cfg.setdefault("device", device)
+ diar_segments = diarize(preprocessed_path, diar_cfg)
+
+ if not diar_segments:
+ logger.warning("No speech segments detected in %s", call_id)
+ return Stage1Output(
+ call_id=call_id,
+ duration=duration,
+ speakers=["speaker_0", "speaker_1"],
+ audio_path=audio_path,
+ segments=[],
+ processing_info=ProcessingInfo(
+ processing_time_sec=time.time() - start_time,
+ models=Models(
+ diarization=diar_cfg.get("model", "pyannote/speaker-diarization-3.1"),
+ asr=cfg.get("asr", {}).get("model", "large-v3-turbo"),
+ language_id=cfg.get("language_id", {}).get("model", "none"),
+ alignment="none",
+ ),
+ device=device,
+ ),
+ )
+
+ # --- Step 3: ASR Transcription ---
+ asr_cfg = cfg.get("asr", {})
+ asr_result = transcribe(preprocessed_path, asr_cfg, device)
+ whisper_language = asr_result.get("language", "ko")
+
+ # --- Step 3.5: Text-based global language decision ---
+ # Whisper's auto-detected language is unreliable on code-switched or
+ # noisy calls. Decide alignment language from the transcript text
+ # instead, applying the same 50% Korean-character majority rule we
+ # use per-segment in language_id.py. Falls back to whisper_language
+ # when the transcript is too short to be statistically meaningful.
+ full_text = " ".join(seg.get("text", "") for seg in asr_result.get("segments", []))
+ align_language, korean_ratio = decide_language_from_text(
+ full_text, fallback=whisper_language,
+ )
+ if align_language != whisper_language:
+ logger.warning(
+ "Language disagreement: whisper=%s, text-based=%s "
+ "(korean_ratio=%.3f, chars=%d)",
+ whisper_language, align_language, korean_ratio,
+ len(full_text.replace(" ", "")),
+ )
+
+ # --- Step 4: Forced Alignment ---
+ align_enabled = cfg.get("alignment", {}).get("enabled", True)
+ aligned_result = align(
+ asr_result, preprocessed_path, device, align_enabled,
+ language_code=align_language,
+ )
+
+ # --- Step 5: Diarization-Transcript Merge ---
+ merged_result = assign_speakers(aligned_result, diar_segments)
+
+ # --- Step 6: Language Detection ---
+ lid_cfg = cfg.get("language_id", {})
+ merged_segments = merged_result.get("segments", [])
+ languages = detect_languages(
+ merged_segments, preprocessed_path, lid_cfg, device, whisper_language,
+ )
+
+ # --- Step 7: Segment Audio Extraction ---
+ segments_dir = cfg.get("segments_dir", "data/segments")
+
+ # Build segment dicts for audio extraction
+ segment_dicts: list[dict] = []
+ for idx, seg in enumerate(merged_segments):
+ speaker = seg.get("speaker", "speaker_0")
+ segment_dicts.append({
+ "segment_id": idx,
+ "speaker_id": speaker,
+ "start": seg.get("start", 0.0),
+ "end": seg.get("end", 0.0),
+ "text": seg.get("text", "").strip(),
+ "language": languages[idx] if idx < len(languages) else whisper_language,
+ })
+
+ audio_paths = _extract_segment_audio(
+ waveform, sr, segment_dicts, call_id, segments_dir,
+ )
+
+ # --- Step 8: Output Assembly ---
+ # Determine speakers list
+ all_speakers = sorted(set(s["speaker_id"] for s in segment_dicts))
+ if len(all_speakers) < 2:
+ all_speakers = ["speaker_0", "speaker_1"]
+
+ segments: list[Segment] = []
+ for seg_dict, seg_audio_path in zip(segment_dicts, audio_paths):
+ # Compute confidence from WhisperX word scores if available
+ orig_seg = merged_segments[seg_dict["segment_id"]] if seg_dict["segment_id"] < len(merged_segments) else {}
+ words = orig_seg.get("words", [])
+ if words and seg_dict["text"]:
+ scores = [w.get("score", 0.0) for w in words if "score" in w]
+ confidence = sum(scores) / len(scores) if scores else 0.5
+ elif not seg_dict["text"]:
+ confidence = 0.0
+ else:
+ confidence = 0.5
+
+ segments.append(Segment(
+ segment_id=seg_dict["segment_id"],
+ speaker_id=seg_dict["speaker_id"],
+ start=round(seg_dict["start"], 3),
+ end=round(seg_dict["end"], 3),
+ text=seg_dict["text"],
+ language=seg_dict["language"],
+ audio_path=seg_audio_path,
+ confidence=round(min(max(confidence, 0.0), 1.0), 3),
+ ))
+
+ processing_time = time.time() - start_time
+
+ output = Stage1Output(
+ call_id=call_id,
+ duration=round(duration, 3),
+ speakers=all_speakers,
+ audio_path=audio_path,
+ segments=segments,
+ processing_info=ProcessingInfo(
+ processing_time_sec=round(processing_time, 3),
+ models=Models(
+ diarization=diar_cfg.get("model", "pyannote/speaker-diarization-3.1"),
+ asr=f"whisperx/{asr_cfg.get('model', 'large-v3-turbo')}-{asr_cfg.get('compute_type', 'int8')}",
+ language_id=lid_cfg.get("model", "none") if lid_cfg.get("enabled", True) else "none",
+ alignment="whisperx/wav2vec2-forced-alignment" if align_enabled else "none",
+ ),
+ device=device,
+ language_chosen=align_language,
+ korean_ratio=round(korean_ratio, 4),
+ ),
+ )
+
+ # Save to JSON
+ output_path = cfg.get("output_path", "data/stage1_output.json")
+ os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
+ with open(output_path, "w", encoding="utf-8") as f:
+ f.write(output.model_dump_json(indent=2))
+
+ logger.info(
+ "Stage 1 complete: %d segments, %.1fs processing time",
+ len(segments), processing_time,
+ )
+
+ return output
diff --git a/src/stage1/requirements.txt b/src/stage1/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..67e31597804aff248fd135be13607462649ca555
--- /dev/null
+++ b/src/stage1/requirements.txt
@@ -0,0 +1,7 @@
+pyannote.audio>=4.0
+faster-whisper>=1.0.0
+whisperx>=3.1.0
+funasr>=1.0.0
+torch>=2.0.0
+torchaudio>=2.0.0
+pyyaml>=6.0
diff --git a/src/stage1/transcription.py b/src/stage1/transcription.py
new file mode 100644
index 0000000000000000000000000000000000000000..b9a37731616fd2c5f85c6c26a37d5284f05557c1
--- /dev/null
+++ b/src/stage1/transcription.py
@@ -0,0 +1,205 @@
+"""Stage 1 โ ASR Transcription Module.
+
+Uses Faster-Whisper (large-v3-turbo, INT8) via WhisperX for
+transcription, forced alignment, and speaker assignment.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import whisperx
+
+logger = logging.getLogger(__name__)
+
+_model = None
+_align_model = None
+_align_metadata = None
+
+
+def _load_model(model_name: str, device: str, compute_type: str):
+ """Load and cache the WhisperX ASR model."""
+ global _model
+ if _model is not None:
+ return _model
+
+ logger.info("Loading WhisperX model: %s (%s)", model_name, compute_type)
+ _model = whisperx.load_model(
+ model_name,
+ device=device,
+ compute_type=compute_type,
+ )
+ return _model
+
+
+def transcribe(audio_path: str, config: dict, device: str) -> dict:
+ """Transcribe full audio using WhisperX.
+
+ Args:
+ audio_path: Path to 16kHz mono WAV file.
+ config: stage1.asr config section.
+ device: Compute device ("cuda" or "cpu").
+
+ Returns:
+ WhisperX result dict with segments and detected language.
+ """
+ model_name = config.get("model", "large-v3-turbo")
+ compute_type = config.get("compute_type", "int8")
+ batch_size = config.get("batch_size", 16)
+ language = config.get("language") # None for auto-detect
+
+ model = _load_model(model_name, device, compute_type)
+
+ logger.info("Transcribing %s (batch_size=%d)", audio_path, batch_size)
+ audio = whisperx.load_audio(audio_path)
+ result = model.transcribe(audio, batch_size=batch_size, language=language)
+
+ detected_lang = result.get("language", "ko")
+ logger.info(
+ "Transcription complete: %d segments, detected language: %s",
+ len(result.get("segments", [])), detected_lang,
+ )
+ return result
+
+
+def align(
+ result: dict,
+ audio_path: str,
+ device: str,
+ align_enabled: bool = True,
+ language_code: str | None = None,
+) -> dict:
+ """Apply forced alignment to get word-level timestamps.
+
+ Args:
+ result: WhisperX transcription result.
+ audio_path: Path to audio file.
+ device: Compute device.
+ align_enabled: Whether to run alignment (from config).
+ language_code: Override for the wav2vec2 model language. When
+ provided, takes precedence over Whisper's auto-detected
+ ``result["language"]``. The orchestrator passes this in
+ after computing language from the transcribed text, which
+ avoids Whisper's silent SPOF on misclassified audio.
+
+ Returns:
+ Updated result dict with word-level timestamps.
+ """
+ if not align_enabled:
+ logger.info("Alignment disabled, skipping")
+ return result
+
+ global _align_model, _align_metadata
+
+ language = language_code if language_code is not None else result.get("language", "ko")
+
+ try:
+ if _align_model is None or _align_metadata is None:
+ logger.info("Loading alignment model for language: %s", language)
+ _align_model, _align_metadata = whisperx.load_align_model(
+ language_code=language, device=device,
+ )
+
+ audio = whisperx.load_audio(audio_path)
+ result = whisperx.align(
+ result["segments"], _align_model, _align_metadata, audio, device,
+ return_char_alignments=False,
+ )
+ logger.info("Alignment complete")
+ except Exception as e:
+ logger.warning("Alignment failed (%s), using Whisper timestamps", e)
+
+ return result
+
+
+def assign_speakers(result: dict, diar_segments: list) -> dict:
+ """Assign speaker labels from diarization to transcription segments.
+
+ After WhisperX word-level speaker assignment, re-segments by speaker
+ boundaries so that multi-speaker ASR segments are properly split.
+
+ Args:
+ result: Aligned WhisperX result.
+ diar_segments: List of DiarSegment from diarization module.
+
+ Returns:
+ Updated result with speaker labels per segment/word,
+ re-segmented at speaker boundaries.
+ """
+ import pandas as pd
+
+ diar_df = pd.DataFrame([
+ {"start": s.start, "end": s.end, "speaker": s.speaker_id}
+ for s in diar_segments
+ ])
+
+ result = whisperx.assign_word_speakers(diar_df, result)
+
+ # Re-segment: split segments at speaker boundaries using word-level labels
+ result["segments"] = _resegment_by_speaker(result.get("segments", []))
+
+ logger.info(
+ "Speaker assignment complete: %d segments",
+ len(result.get("segments", [])),
+ )
+ return result
+
+
+def _resegment_by_speaker(segments: list[dict]) -> list[dict]:
+ """Split segments that contain multiple speakers into per-speaker segments.
+
+ WhisperX assigns a single speaker per segment based on majority overlap,
+ but individual words may have different speaker labels. This function
+ groups consecutive same-speaker words into new segments.
+ """
+ new_segments = []
+
+ for seg in segments:
+ words = seg.get("words", [])
+ if not words:
+ new_segments.append(seg)
+ continue
+
+ # Group consecutive words by speaker
+ groups = []
+ current_speaker = None
+ current_words = []
+
+ for word in words:
+ word_speaker = word.get("speaker", seg.get("speaker", "speaker_0"))
+ if word_speaker != current_speaker:
+ if current_words:
+ groups.append((current_speaker, current_words))
+ current_speaker = word_speaker
+ current_words = [word]
+ else:
+ current_words.append(word)
+
+ if current_words:
+ groups.append((current_speaker, current_words))
+
+ # If only one speaker group, keep the original segment
+ if len(groups) <= 1:
+ new_segments.append(seg)
+ continue
+
+ # Create new segments for each speaker group
+ for speaker, group_words in groups:
+ starts = [w.get("start") for w in group_words if w.get("start") is not None]
+ ends = [w.get("end") for w in group_words if w.get("end") is not None]
+ if not starts or not ends:
+ continue
+
+ text = " ".join(w.get("word", "") for w in group_words).strip()
+ if not text:
+ continue
+
+ new_segments.append({
+ "start": min(starts),
+ "end": max(ends),
+ "text": text,
+ "speaker": speaker,
+ "words": group_words,
+ })
+
+ return new_segments
diff --git a/src/stage2/__init__.py b/src/stage2/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..cb8683e38f82f3e1b5cdc04002da8cd383fc5adb
--- /dev/null
+++ b/src/stage2/__init__.py
@@ -0,0 +1,5 @@
+"""Stage 2: Audio Emotion + Text Emotion Analysis."""
+
+from src.stage2.process import process
+
+__all__ = ["process"]
diff --git a/src/stage2/audio_emotion.py b/src/stage2/audio_emotion.py
new file mode 100644
index 0000000000000000000000000000000000000000..c0b190d1ba5b3690be463a03f623c99eb3c857a7
--- /dev/null
+++ b/src/stage2/audio_emotion.py
@@ -0,0 +1,286 @@
+"""Stage 2 โ Audio Emotion Recognition Module.
+
+Uses emotion2vec_plus_base via FunASR for per-segment audio emotion classification.
+Outputs 7-class emotions mapped from emotion2vec's native 9-class taxonomy.
+
+Supports both zero-shot (9โ7 mapping) and fine-tuned (6-class + disgust=0) modes.
+Set `finetuned_checkpoint` in config.yaml to switch to fine-tuned mode.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+# emotion2vec 9-class โ project 7-class mapping
+# emotion2vec outputs labels like "็ๆฐ/angry" (Chinese/English) or just "angry"
+LABEL_MAP = {
+ "angry": "anger",
+ "disgusted": "disgust",
+ "fearful": "fear",
+ "happy": "joy",
+ "neutral": "neutral",
+ "sad": "sadness",
+ "surprised": "surprise",
+ "other": "neutral",
+ "unknown": "neutral",
+ # Chinese/English composite labels from emotion2vec
+ "็ๆฐ/angry": "anger",
+ "ๅๆถ/disgusted": "disgust",
+ "ๆๆง/fearful": "fear",
+ "ๅผๅฟ/happy": "joy",
+ "ไธญ็ซ/neutral": "neutral",
+ "้พ่ฟ/sad": "sadness",
+ "ๅๆ/surprised": "surprise",
+ "ๅ
ถไป/other": "neutral",
+ "": "neutral",
+}
+
+PROJECT_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear", "disgust"]
+FINETUNE_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear"]
+
+# LoRA 7-class labels โ project labels (happiness โ joy)
+LORA_LABELS = ["happiness", "anger", "disgust", "fear", "neutral", "sadness", "surprise"]
+LORA_TO_PROJECT = {
+ "happiness": "joy", "anger": "anger", "disgust": "disgust",
+ "fear": "fear", "neutral": "neutral", "sadness": "sadness", "surprise": "surprise",
+}
+
+_model = None
+_finetuned_encoder = None
+_onnx_session = None
+
+
+def _load_model(model_id: str, device: str):
+ """Load and cache the emotion2vec model (singleton)."""
+ global _model
+ if _model is not None:
+ return _model
+
+ try:
+ from funasr import AutoModel
+ except ImportError:
+ raise RuntimeError(
+ "funasr is required for audio emotion recognition. "
+ "Install with: pip install funasr onnxruntime"
+ )
+
+ logger.info("Loading audio emotion model: %s", model_id)
+ _model = AutoModel(model=model_id, device=device, hub="hf")
+ logger.info("Audio emotion model loaded on %s", device)
+ return _model
+
+
+def _load_finetuned(model_id: str, checkpoint_path: str, device: str):
+ """Load fine-tuned emotion2vec with custom 6-class proj head."""
+ global _finetuned_encoder
+ if _finetuned_encoder is not None:
+ return _finetuned_encoder
+
+ fmodel = _load_model(model_id, device)
+ encoder = fmodel.model
+
+ # Replace proj with 6-class head
+ encoder.proj = torch.nn.Linear(768, len(FINETUNE_LABELS)).to(device)
+
+ # Load fine-tuned weights
+ ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+ encoder.blocks.load_state_dict(ckpt["blocks"])
+ encoder.proj.load_state_dict(ckpt["proj"])
+ if "norm" in ckpt and encoder.norm is not None:
+ encoder.norm.load_state_dict(ckpt["norm"])
+
+ encoder.eval()
+ _finetuned_encoder = encoder
+ logger.info("Loaded fine-tuned checkpoint: %s", checkpoint_path)
+ return encoder
+
+
+def _map_scores(raw_labels: list[str], raw_scores: list[float]) -> dict[str, float]:
+ """Map emotion2vec native labels to project 7-class taxonomy."""
+ mapped = {label: 0.0 for label in PROJECT_LABELS}
+ for native_label, score in zip(raw_labels, raw_scores):
+ project_label = LABEL_MAP.get(native_label, "neutral")
+ mapped[project_label] += float(score)
+
+ # Normalize to sum to 1.0
+ total = sum(mapped.values())
+ if total > 0:
+ mapped = {k: v / total for k, v in mapped.items()}
+
+ return mapped
+
+
+def _load_onnx(onnx_path: str):
+ """Load and cache the LoRA ONNX model (singleton)."""
+ global _onnx_session
+ if _onnx_session is not None:
+ return _onnx_session
+
+ import onnxruntime as ort
+ logger.info("Loading LoRA ONNX model: %s", onnx_path)
+ _onnx_session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
+ logger.info("LoRA ONNX model loaded")
+ return _onnx_session
+
+
+def _predict_lora_onnx(audio_path: str, session) -> dict[str, float]:
+ """Inference with LoRA ONNX model (7-class, happinessโjoy mapping)."""
+ import numpy as np
+ import soundfile as sf
+
+ audio, sr = sf.read(audio_path, dtype="float32")
+ if audio.ndim == 2:
+ audio = audio.mean(axis=1)
+ if sr != 16000:
+ import librosa
+ audio = librosa.resample(audio, orig_sr=sr, target_sr=16000)
+
+ waveform = audio.reshape(1, -1).astype(np.float32)
+ logits = session.run(None, {"waveform": waveform})[0]
+
+ # Softmax
+ exp_logits = np.exp(logits - logits.max(axis=-1, keepdims=True))
+ probs = (exp_logits / exp_logits.sum(axis=-1, keepdims=True)).squeeze()
+
+ # Map LoRA labels โ project labels (happiness โ joy)
+ scores = {label: 0.0 for label in PROJECT_LABELS}
+ for lora_label, prob in zip(LORA_LABELS, probs):
+ project_label = LORA_TO_PROJECT[lora_label]
+ scores[project_label] = float(prob)
+
+ return scores
+
+
+def _predict_finetuned(audio_path: str, encoder, device: str) -> dict[str, float]:
+ """Inference with fine-tuned model (direct forward pass, 6-class)."""
+ import soundfile as sf
+
+ audio, sr = sf.read(audio_path, dtype="float32")
+ if audio.ndim == 2:
+ audio = audio.mean(axis=1)
+ if sr != 16000:
+ import librosa
+ audio = librosa.resample(audio, orig_sr=sr, target_sr=16000)
+
+ source = torch.tensor(audio, dtype=torch.float32).to(device)
+ if encoder.cfg.normalize:
+ source = F.layer_norm(source, source.shape)
+ source = source.unsqueeze(0)
+
+ with torch.no_grad():
+ feats = encoder.extract_features(source, padding_mask=None)
+ pooled = feats["x"].mean(dim=1)
+ logits = encoder.proj(pooled)
+ probs = torch.softmax(logits, dim=-1).squeeze().cpu().tolist()
+
+ # 6-class โ 7-class (add disgust=0.0 for compatibility)
+ scores = {label: 0.0 for label in PROJECT_LABELS}
+ for label, prob in zip(FINETUNE_LABELS, probs):
+ scores[label] = float(prob)
+
+ # Normalize
+ total = sum(scores.values())
+ if total > 0:
+ scores = {k: v / total for k, v in scores.items()}
+ return scores
+
+
+def predict(
+ audio_path: str,
+ device: str = "cpu",
+ model_id: str = "iic/emotion2vec_plus_base",
+ finetuned_checkpoint: str | None = None,
+ lora_onnx_path: str | None = None,
+) -> dict:
+ """Predict emotion from an audio segment.
+
+ Args:
+ audio_path: Path to 16kHz mono WAV segment file.
+ device: Compute device ("cpu" or "cuda").
+ model_id: HuggingFace model ID for emotion2vec.
+ finetuned_checkpoint: Path to fine-tuned checkpoint (None = zero-shot).
+ lora_onnx_path: Path to LoRA ONNX model (takes priority over other modes).
+
+ Returns:
+ {"emotion": str, "confidence": float, "scores": dict[str, float]}
+ """
+ # Check file exists
+ if not Path(audio_path).exists():
+ logger.warning("Audio file not found: %s โ returning neutral", audio_path)
+ return {
+ "emotion": "neutral",
+ "confidence": 0.0,
+ "scores": {label: 1.0 / len(PROJECT_LABELS) for label in PROJECT_LABELS},
+ }
+
+ # LoRA ONNX path (highest priority)
+ if lora_onnx_path and Path(lora_onnx_path).exists():
+ session = _load_onnx(lora_onnx_path)
+ try:
+ scores = _predict_lora_onnx(audio_path, session)
+ except Exception as e:
+ logger.warning("LoRA ONNX inference failed on %s: %s โ returning neutral", audio_path, e)
+ scores = {label: 1.0 / len(PROJECT_LABELS) for label in PROJECT_LABELS}
+
+ top_label = max(scores, key=scores.get)
+ return {
+ "emotion": top_label,
+ "confidence": round(scores[top_label], 4),
+ "scores": {k: round(v, 4) for k, v in scores.items()},
+ }
+
+ # Fine-tuned path
+ if finetuned_checkpoint and Path(finetuned_checkpoint).exists():
+ encoder = _load_finetuned(model_id, finetuned_checkpoint, device)
+ try:
+ scores = _predict_finetuned(audio_path, encoder, device)
+ except Exception as e:
+ logger.warning("Fine-tuned inference failed on %s: %s โ returning neutral", audio_path, e)
+ scores = {label: 1.0 / len(PROJECT_LABELS) for label in PROJECT_LABELS}
+
+ top_label = max(scores, key=scores.get)
+ return {
+ "emotion": top_label,
+ "confidence": round(scores[top_label], 4),
+ "scores": {k: round(v, 4) for k, v in scores.items()},
+ }
+
+ # Zero-shot path (original)
+ model = _load_model(model_id, device)
+
+ try:
+ output = model.generate(
+ audio_path, granularity="utterance", extract_embedding=False,
+ )
+ except Exception as e:
+ logger.warning("Audio emotion inference failed on %s: %s โ returning neutral", audio_path, e)
+ return {
+ "emotion": "neutral",
+ "confidence": 0.0,
+ "scores": {label: 1.0 / len(PROJECT_LABELS) for label in PROJECT_LABELS},
+ }
+
+ # Parse emotion2vec output
+ if output and isinstance(output, list) and len(output) > 0:
+ rec = output[0]
+ raw_labels = rec.get("labels", [])
+ raw_scores = rec.get("scores", [])
+ scores = _map_scores(raw_labels, raw_scores)
+ else:
+ logger.warning("Empty model output for %s โ returning neutral", audio_path)
+ scores = {label: 1.0 / len(PROJECT_LABELS) for label in PROJECT_LABELS}
+
+ top_label = max(scores, key=scores.get)
+ confidence = scores[top_label]
+
+ return {
+ "emotion": top_label,
+ "confidence": round(confidence, 4),
+ "scores": {k: round(v, 4) for k, v in scores.items()},
+ }
diff --git a/src/stage2/features.py b/src/stage2/features.py
new file mode 100644
index 0000000000000000000000000000000000000000..66145d07bf32f87264a2da9176a937e5a63cdfd3
--- /dev/null
+++ b/src/stage2/features.py
@@ -0,0 +1,17 @@
+"""Stage 2 โ Audio Feature Extraction (Placeholder).
+
+Will be implemented with librosa/openSMILE during fine-tuning phase.
+Currently returns empty dict to maintain interface compatibility.
+"""
+
+
+def extract_features(audio_path: str) -> dict:
+ """Placeholder โ extract audio features for future use.
+
+ Args:
+ audio_path: Path to 16kHz mono WAV segment file.
+
+ Returns:
+ Empty dict (placeholder for librosa/openSMILE features).
+ """
+ return {}
diff --git a/src/stage2/fusion.py b/src/stage2/fusion.py
new file mode 100644
index 0000000000000000000000000000000000000000..3997d0b3690eca0146e26fd985cdbd2cbef69b22
--- /dev/null
+++ b/src/stage2/fusion.py
@@ -0,0 +1,76 @@
+"""Stage 2 โ Audio + Text Emotion Fusion.
+
+Uses emotion-specific fusion weights per class (from Korean 263 val grid search).
+Falls back to fixed 60/40 if mode="fixed".
+"""
+
+from __future__ import annotations
+
+from src.common.constants import (
+ EMOTION_FUSION_WEIGHTS_EN,
+ EMOTION_FUSION_WEIGHTS_KO,
+ FUSION_WEIGHTS,
+)
+
+PROJECT_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear", "disgust"]
+
+
+def fuse(
+ audio_scores: dict[str, float],
+ text_scores: dict[str, float],
+ mode: str = "emotion_specific",
+ language: str = "ko",
+) -> dict:
+ """Fuse audio and text emotion predictions.
+
+ Args:
+ audio_scores: {"emotion": str, "confidence": float, "scores": dict}
+ text_scores: {"emotion": str, "confidence": float, "scores": dict}
+ If text was empty, scores will be uniform distribution.
+ mode: "emotion_specific" (per-class weights) or "fixed" (60/40)
+ language: "ko" or "en" โ selects language-specific emotion weights.
+
+ Returns:
+ {"emotion": str, "confidence": float, "scores": dict[str, float]}
+ """
+ a_scores = audio_scores.get("scores", {})
+ t_scores = text_scores.get("scores", {})
+
+ # If text confidence is very low (empty text), use audio only
+ text_confidence = text_scores.get("confidence", 0.0)
+ audio_only = text_confidence <= 0.1
+
+ weights_table = EMOTION_FUSION_WEIGHTS_EN if language == "en" else EMOTION_FUSION_WEIGHTS_KO
+
+ # Weighted average of score distributions
+ fused = {}
+ for label in PROJECT_LABELS:
+ a = a_scores.get(label, 0.0)
+ t = t_scores.get(label, 0.0)
+
+ if audio_only:
+ audio_w, text_w = 1.0, 0.0
+ elif mode == "emotion_specific":
+ w = weights_table.get(label, {"audio": 0.6, "text": 0.4})
+ audio_w, text_w = w["audio"], w["text"]
+ else:
+ audio_w = FUSION_WEIGHTS["audio"]
+ text_w = FUSION_WEIGHTS["text"]
+
+ fused[label] = a * audio_w + t * text_w
+
+ # Normalize to sum to 1.0
+ total = sum(fused.values())
+ if total > 0:
+ fused = {k: v / total for k, v in fused.items()}
+ else:
+ fused = {label: 1.0 / len(PROJECT_LABELS) for label in PROJECT_LABELS}
+
+ top_label = max(fused, key=fused.get)
+ confidence = fused[top_label]
+
+ return {
+ "emotion": top_label,
+ "confidence": round(confidence, 4),
+ "scores": {k: round(v, 4) for k, v in fused.items()},
+ }
diff --git a/src/stage2/process.py b/src/stage2/process.py
new file mode 100644
index 0000000000000000000000000000000000000000..0c2cb30dfc6ca8abde97e04e16a3ab625f79dafb
--- /dev/null
+++ b/src/stage2/process.py
@@ -0,0 +1,204 @@
+"""Stage 2 โ Orchestrator.
+
+Entry point: process(stage1_output) -> Stage2Output
+
+Pipeline:
+ 1. Load config
+ 2. Per segment: audio emotion + text emotion + fusion
+ 3. Aggregate per speaker โ SpeakerSummary
+ 4. Assemble Stage2Output + save JSON
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import time
+from collections import defaultdict
+from pathlib import Path
+
+import torch
+import yaml
+
+from src.common.constants import EMOTION_LABELS
+from src.common.schemas import (
+ EmotionResult,
+ SpeakerSummary,
+ Stage1Output,
+ Stage2Output,
+)
+from src.stage2.audio_emotion import predict as audio_predict
+from src.stage2.fusion import fuse
+from src.stage2.text_emotion import predict as text_predict
+
+logger = logging.getLogger(__name__)
+
+
+def _load_config(config: dict | None = None) -> dict:
+ """Load config from yaml if not provided."""
+ if config is not None:
+ return config
+
+ config_path = Path(__file__).parent.parent.parent / "config.yaml"
+ if config_path.exists():
+ with open(config_path) as f:
+ return yaml.safe_load(f).get("stage2", {})
+
+ logger.warning("config.yaml not found, using defaults")
+ return {}
+
+
+def _get_device() -> str:
+ """Determine compute device."""
+ return "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def _aggregate_speakers(
+ emotions: list[EmotionResult],
+) -> dict[str, SpeakerSummary]:
+ """Aggregate per-segment emotions into per-speaker summaries.
+
+ For each speaker:
+ - dominant_emotion: emotion with highest total weight
+ - emotion_distribution: normalized weights (sum = 1.0)
+ - avg_confidence: mean fused_confidence
+ """
+ speaker_emotions: dict[str, list[EmotionResult]] = defaultdict(list)
+ for em in emotions:
+ speaker_emotions[em.speaker_id].append(em)
+
+ summaries: dict[str, SpeakerSummary] = {}
+ for speaker_id, speaker_ems in speaker_emotions.items():
+ # Count emotion occurrences weighted by confidence
+ emotion_weights: dict[str, float] = defaultdict(float)
+ total_confidence = 0.0
+
+ for em in speaker_ems:
+ emotion_weights[em.fused_emotion] += em.fused_confidence
+ total_confidence += em.fused_confidence
+
+ # Normalize to sum to 1.0
+ total_weight = sum(emotion_weights.values())
+ if total_weight > 0:
+ emotion_distribution = {
+ k: round(v / total_weight, 4)
+ for k, v in emotion_weights.items()
+ }
+ else:
+ emotion_distribution = {"neutral": 1.0}
+
+ dominant_emotion = max(emotion_distribution, key=emotion_distribution.get)
+ avg_confidence = total_confidence / len(speaker_ems) if speaker_ems else 0.0
+
+ summaries[speaker_id] = SpeakerSummary(
+ dominant_emotion=dominant_emotion,
+ emotion_distribution=emotion_distribution,
+ avg_confidence=round(avg_confidence, 4),
+ )
+
+ return summaries
+
+
+def process(
+ stage1_output: Stage1Output,
+ config: dict | None = None,
+) -> Stage2Output:
+ """Stage 2 entry point.
+
+ Args:
+ stage1_output: Stage 1 output with segments containing audio paths and text.
+ config: Optional config dict. If None, loads from config.yaml.
+
+ Returns:
+ Stage2Output with per-segment emotions and per-speaker summaries.
+ """
+ start_time = time.time()
+ cfg = _load_config(config)
+ device = _get_device()
+
+ audio_cfg = cfg.get("audio_emotion", {})
+ audio_model_id = audio_cfg.get("model", "iic/emotion2vec_plus_base")
+ lora_onnx_path = audio_cfg.get("lora_onnx_path")
+
+ text_cfg = cfg.get("text_emotion", {})
+ ko_model_id = text_cfg.get("korean_model")
+ en_model_id = text_cfg.get("english_model")
+ ko_lora_onnx = text_cfg.get("korean_lora_onnx_path")
+ ko_lora_tokenizer = text_cfg.get("korean_lora_tokenizer")
+
+ fusion_cfg = cfg.get("fusion", {})
+ fusion_mode = fusion_cfg.get("mode", "emotion_specific")
+
+ call_id = stage1_output.call_id
+ segments = stage1_output.segments
+
+ logger.info(
+ "Stage 2 processing %s: %d segments, device=%s",
+ call_id, len(segments), device,
+ )
+
+ if not segments:
+ logger.warning("No segments to process for %s", call_id)
+ return Stage2Output(
+ call_id=call_id,
+ emotions=[],
+ speaker_summaries={},
+ )
+
+ # Per-segment emotion analysis
+ emotions: list[EmotionResult] = []
+ for i, seg in enumerate(segments):
+ # Audio emotion (LoRA ONNX preferred if available)
+ audio_result = audio_predict(
+ seg.audio_path, device=device, model_id=audio_model_id,
+ lora_onnx_path=lora_onnx_path,
+ )
+
+ # Text emotion (Korean uses LoRA ONNX if available)
+ text_model_id = ko_model_id if seg.language == "ko" else en_model_id
+ text_result = text_predict(
+ seg.text, language=seg.language, model_id=text_model_id,
+ korean_lora_onnx_path=ko_lora_onnx,
+ korean_lora_tokenizer=ko_lora_tokenizer,
+ )
+
+ # Fusion (emotion-specific by default, per-language weights)
+ fused_result = fuse(audio_result, text_result, mode=fusion_mode, language=seg.language)
+
+ emotions.append(EmotionResult(
+ speaker_id=seg.speaker_id,
+ segment_id=seg.segment_id,
+ audio_emotion=audio_result["emotion"],
+ audio_confidence=audio_result["confidence"],
+ text_emotion=text_result["emotion"],
+ text_confidence=text_result["confidence"],
+ fused_emotion=fused_result["emotion"],
+ fused_confidence=fused_result["confidence"],
+ ))
+
+ if (i + 1) % 10 == 0:
+ logger.info(" %d/%d segments processed", i + 1, len(segments))
+
+ # Speaker aggregation
+ speaker_summaries = _aggregate_speakers(emotions)
+
+ processing_time = time.time() - start_time
+
+ output = Stage2Output(
+ call_id=call_id,
+ emotions=emotions,
+ speaker_summaries=speaker_summaries,
+ )
+
+ # Save to JSON
+ output_path = cfg.get("output_path", "data/stage2_output.json")
+ os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
+ with open(output_path, "w", encoding="utf-8") as f:
+ f.write(output.model_dump_json(indent=2))
+
+ logger.info(
+ "Stage 2 complete: %d emotions, %d speakers, %.1fs processing time",
+ len(emotions), len(speaker_summaries), processing_time,
+ )
+
+ return output
diff --git a/src/stage2/requirements.txt b/src/stage2/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..b61066aec4ccba7068cf8ac55c2a3dcd9a03051a
--- /dev/null
+++ b/src/stage2/requirements.txt
@@ -0,0 +1,6 @@
+onnxruntime>=1.17.0
+transformers>=4.38.0
+torch>=2.0.0
+librosa>=0.10.0
+opensmile>=2.5.0
+numpy>=1.24.0
diff --git a/src/stage2/text_emotion.py b/src/stage2/text_emotion.py
new file mode 100644
index 0000000000000000000000000000000000000000..e97791d6b7f50517caf8575b9b58450e8434a20a
--- /dev/null
+++ b/src/stage2/text_emotion.py
@@ -0,0 +1,260 @@
+"""Stage 2 โ Text Emotion Recognition Module.
+
+Uses KcELECTRA (Korean) and DistilRoBERTa (English) via transformers pipeline
+for per-segment text emotion classification.
+
+If a LoRA-finetuned KcELECTRA ONNX is provided (Korean), it takes priority.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+PROJECT_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear", "disgust"]
+
+# LoRA KcELECTRA 7-class labels โ project labels (happiness โ joy)
+LORA_KO_LABELS = ["happiness", "anger", "disgust", "fear", "neutral", "sadness", "surprise"]
+LORA_KO_TO_PROJECT = {
+ "happiness": "joy", "anger": "anger", "disgust": "disgust",
+ "fear": "fear", "neutral": "neutral", "sadness": "sadness", "surprise": "surprise",
+}
+
+# Label mappings from model-specific labels to project taxonomy
+
+# j-hartmann/emotion-english-distilroberta-base: 7 emotion labels
+EN_LABEL_MAP = {
+ "anger": "anger",
+ "disgust": "disgust",
+ "fear": "fear",
+ "joy": "joy",
+ "neutral": "neutral",
+ "sadness": "sadness",
+ "surprise": "surprise",
+}
+
+# searle-j/kote_for_easygoing_people: 44 fine-grained Korean emotions โ 7 project labels
+KO_LABEL_MAP = {
+ # joy
+ "๊ธฐ์จ": "joy",
+ "์ฆ๊ฑฐ์/์ ๋จ": "joy",
+ "ํ๋ณต": "joy",
+ "๊ฐ๋/๊ฐํ": "joy",
+ "๊ณ ๋ง์": "joy",
+ "ํ์/ํธ์": "joy",
+ "๋ฟ๋ฏํจ": "joy",
+ "ํ๋ญํจ(๊ท์ฌ์/์์จ)": "joy",
+ "๊ธฐ๋๊ฐ": "joy",
+ "ํธ์/์พ์ ": "joy",
+ "์์ฌ/์ ๋ขฐ": "joy",
+ "์๊ปด์ฃผ๋": "joy",
+ "์กด๊ฒฝ": "joy",
+ # surprise
+ "๋๋": "surprise",
+ "์ ๊ธฐํจ/๊ด์ฌ": "surprise",
+ "๊ฒฝ์
": "surprise",
+ # sadness
+ "์ฌํ": "sadness",
+ "์๋ฌ์": "sadness",
+ "์ํ๊น์/์ค๋ง": "sadness",
+ "๋ถ์ํจ/์ฐ๋ฏผ": "sadness",
+ "ํ๋ฆ/์ง์นจ": "sadness",
+ "์ ๋ง": "sadness",
+ "๋น์ฅํจ": "sadness",
+ "ํจ๋ฐฐ/์๊ธฐํ์ค": "sadness",
+ # anger
+ "ํ๋จ/๋ถ๋
ธ": "anger",
+ "์ง์ฆ": "anger",
+ "๋ถํ/๋ถ๋ง": "anger",
+ "์ง๊ธ์ง๊ธ": "anger",
+ "ํ์ฌํจ": "anger",
+ "์ฐ์ญ๋/๋ฌด์ํจ": "anger",
+ # fear
+ "๋ถ์/๊ฑฑ์ ": "fear",
+ "๊ณตํฌ/๋ฌด์์": "fear",
+ "๋ถ๋ด/์_๋ดํด": "fear",
+ # disgust
+ "์ฆ์ค/ํ์ค": "disgust",
+ "์ญ๊ฒจ์/์ง๊ทธ๋ฌ์": "disgust",
+ "์์ฌ/๋ถ์ ": "disgust",
+ # neutral
+ "์์": "neutral",
+ "๊นจ๋ฌ์": "neutral",
+ "์ฌ๋ฏธ์์": "neutral",
+ "๋นํฉ/๋์ฒ": "neutral",
+ "์ด์ด์์": "neutral",
+ "๋ถ๋๋ฌ์": "neutral",
+ "๊ท์ฐฎ์": "neutral",
+ "์ฃ์ฑ
๊ฐ": "neutral",
+}
+
+_model_ko = None
+_model_en = None
+_onnx_session_ko = None
+_tokenizer_ko = None
+
+
+def _load_lora_ko_onnx(onnx_path: str, tokenizer_path: str):
+ """Load cached LoRA KcELECTRA ONNX session + tokenizer (singleton)."""
+ global _onnx_session_ko, _tokenizer_ko
+ if _onnx_session_ko is not None:
+ return _onnx_session_ko, _tokenizer_ko
+
+ import onnxruntime as ort
+ from transformers import AutoTokenizer
+
+ logger.info("Loading LoRA KcELECTRA ONNX: %s", onnx_path)
+ _onnx_session_ko = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
+ _tokenizer_ko = AutoTokenizer.from_pretrained(tokenizer_path)
+ logger.info("LoRA KcELECTRA ONNX loaded")
+ return _onnx_session_ko, _tokenizer_ko
+
+
+def _predict_lora_ko_onnx(text: str, session, tokenizer) -> dict[str, float]:
+ """7-class direct prediction via LoRA KcELECTRA ONNX."""
+ import numpy as np
+
+ enc = tokenizer(text, return_tensors="np", truncation=True, max_length=128, padding="max_length")
+ logits = session.run(None, {
+ "input_ids": enc["input_ids"],
+ "attention_mask": enc["attention_mask"],
+ })[0]
+
+ exp_logits = np.exp(logits - logits.max(axis=-1, keepdims=True))
+ probs = (exp_logits / exp_logits.sum(axis=-1, keepdims=True)).squeeze()
+
+ scores = {label: 0.0 for label in PROJECT_LABELS}
+ for lora_label, prob in zip(LORA_KO_LABELS, probs):
+ scores[LORA_KO_TO_PROJECT[lora_label]] = float(prob)
+ return scores
+
+
+def _load_model(language: str, model_id: str | None = None):
+ """Load and cache the text emotion model (singleton per language)."""
+ global _model_ko, _model_en
+
+ if language == "ko" and _model_ko is not None:
+ return _model_ko
+ if language != "ko" and _model_en is not None:
+ return _model_en
+
+ try:
+ from transformers import pipeline
+ except ImportError:
+ raise RuntimeError(
+ "transformers is required for text emotion recognition. "
+ "Install with: pip install transformers torch"
+ )
+
+ if language == "ko":
+ model_name = model_id or "searle-j/kote_for_easygoing_people"
+ logger.info("Loading Korean text emotion model: %s", model_name)
+ _model_ko = pipeline(
+ "text-classification", model=model_name, top_k=None, truncation=True,
+ )
+ logger.info("Korean text emotion model loaded")
+ return _model_ko
+ else:
+ model_name = model_id or "j-hartmann/emotion-english-distilroberta-base"
+ logger.info("Loading English text emotion model: %s", model_name)
+ _model_en = pipeline(
+ "text-classification", model=model_name, top_k=None, truncation=True,
+ )
+ logger.info("English text emotion model loaded")
+ return _model_en
+
+
+def _map_scores(raw_results: list[dict], language: str) -> dict[str, float]:
+ """Map model output labels to project 7-class taxonomy."""
+ label_map = KO_LABEL_MAP if language == "ko" else EN_LABEL_MAP
+ mapped = {label: 0.0 for label in PROJECT_LABELS}
+
+ for item in raw_results:
+ native_label = item.get("label", "").lower().strip()
+ score = float(item.get("score", 0.0))
+ project_label = label_map.get(native_label, None)
+ if project_label:
+ mapped[project_label] += score
+ else:
+ # Unknown label โ distribute to neutral
+ mapped["neutral"] += score
+
+ # Normalize
+ total = sum(mapped.values())
+ if total > 0:
+ mapped = {k: v / total for k, v in mapped.items()}
+
+ return mapped
+
+
+def predict(
+ text: str,
+ language: str = "ko",
+ model_id: str | None = None,
+ korean_lora_onnx_path: str | None = None,
+ korean_lora_tokenizer: str | None = None,
+) -> dict:
+ """Predict emotion from text.
+
+ Args:
+ text: Transcribed text from a speech segment.
+ language: Language code ("ko" or "en").
+ model_id: Optional override for model ID.
+ korean_lora_onnx_path: Optional LoRA ONNX path for Korean (takes priority).
+ korean_lora_tokenizer: Tokenizer path/id for Korean LoRA ONNX.
+
+ Returns:
+ {"emotion": str, "confidence": float, "scores": dict[str, float]}
+ """
+ # Handle empty text
+ if not text or not text.strip():
+ return {
+ "emotion": "neutral",
+ "confidence": 0.1,
+ "scores": {label: 1.0 / len(PROJECT_LABELS) for label in PROJECT_LABELS},
+ }
+
+ # LoRA KcELECTRA ONNX path (Korean only, highest priority)
+ if (language == "ko" and korean_lora_onnx_path
+ and Path(korean_lora_onnx_path).exists()
+ and korean_lora_tokenizer and Path(korean_lora_tokenizer).exists()):
+ try:
+ session, tokenizer = _load_lora_ko_onnx(korean_lora_onnx_path, korean_lora_tokenizer)
+ scores = _predict_lora_ko_onnx(text, session, tokenizer)
+ top_label = max(scores, key=scores.get)
+ return {
+ "emotion": top_label,
+ "confidence": round(scores[top_label], 4),
+ "scores": {k: round(v, 4) for k, v in scores.items()},
+ }
+ except Exception as e:
+ logger.warning("LoRA KcELECTRA ONNX failed: %s โ falling back to base model", e)
+
+ model = _load_model(language, model_id)
+
+ try:
+ results = model(text)
+ except Exception as e:
+ logger.warning("Text emotion inference failed: %s โ returning neutral", e)
+ return {
+ "emotion": "neutral",
+ "confidence": 0.1,
+ "scores": {label: 1.0 / len(PROJECT_LABELS) for label in PROJECT_LABELS},
+ }
+
+ # transformers pipeline with top_k=None returns list of dicts
+ if isinstance(results, list) and results and isinstance(results[0], list):
+ # Some pipelines return nested lists
+ results = results[0]
+
+ scores = _map_scores(results, language)
+ top_label = max(scores, key=scores.get)
+ confidence = scores[top_label]
+
+ return {
+ "emotion": top_label,
+ "confidence": round(confidence, 4),
+ "scores": {k: round(v, 4) for k, v in scores.items()},
+ }
diff --git a/src/stage3/__init__.py b/src/stage3/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/src/stage3/character_mapping.py b/src/stage3/character_mapping.py
new file mode 100644
index 0000000000000000000000000000000000000000..d50e5b6c6a3cf286256f2f43810f16f9fab270b9
--- /dev/null
+++ b/src/stage3/character_mapping.py
@@ -0,0 +1,190 @@
+"""
+๊ฐ์ โ ์บ๋ฆญํฐ ๋ฐ์ ๋งคํ (Stage 3)
+
+๋งคํ ์คํ ์ค๊ณ: ์ฃผํ (PO)
+๊ตฌํ: ์๊ท (App Engineer)
+
+solo_state: ๊ฐ ํ์์ ๊ฐ๋ณ ๊ฐ์ (7 Ekman, other/unknown โ neutral)
+pair_state: ๋ ํ์์ ๊ฐ์ ์กฐํฉ์ผ๋ก ๊ฒฐ์ ๋๋ ์ํธ์์ฉ ํฌ์ฆ
+"""
+
+from src.common.schemas import CharacterReaction, SpeakerSummary
+
+# ----- 4-class mood atom (up/calm/down/tense) -----
+
+_DEFAULT_MOOD = {
+ "joy": "up",
+ "neutral": "calm",
+ "sadness": "down",
+ "anger": "tense",
+}
+
+_AMBIGUOUS = {"surprise", "fear", "disgust"}
+
+_CONFIDENCE_GATE = 0.5
+
+
+def _resolve_ambiguous(emotion: str, dist: dict[str, float]) -> str:
+ """Resolve surprise/fear/disgust using the non-dominant residual distribution."""
+ R_up = dist.get("joy", 0.0)
+
+ if emotion == "surprise":
+ R_down = dist.get("sadness", 0.0) + dist.get("disgust", 0.0)
+ R_tense = dist.get("anger", 0.0) + dist.get("fear", 0.0)
+ if R_up > R_tense + R_down:
+ return "up"
+ if R_tense > R_up:
+ return "tense"
+ return "calm"
+
+ if emotion == "fear":
+ R_down = dist.get("sadness", 0.0) + dist.get("disgust", 0.0)
+ R_tense = dist.get("anger", 0.0)
+ if R_tense > R_down:
+ return "tense"
+ if R_down > R_tense:
+ return "down"
+ return "tense"
+
+ if emotion == "disgust":
+ R_down = dist.get("sadness", 0.0)
+ R_tense = dist.get("anger", 0.0) + dist.get("fear", 0.0)
+ if R_tense > R_down:
+ return "tense"
+ if R_down > R_tense:
+ return "down"
+ return "tense"
+
+ return "calm"
+
+
+def mood(summary: SpeakerSummary) -> str:
+ """Reduce a SpeakerSummary to one of {up, calm, down, tense}."""
+ dist = summary.emotion_distribution or {}
+ if not dist:
+ return "calm"
+
+ if max(dist.values(), default=0.0) < _CONFIDENCE_GATE:
+ return "calm"
+
+ dom = (summary.dominant_emotion or "").lower().strip()
+
+ if dom in _DEFAULT_MOOD:
+ return _DEFAULT_MOOD[dom]
+
+ if dom in _AMBIGUOUS:
+ return _resolve_ambiguous(dom, dist)
+
+ return "calm"
+
+
+# ----- 4ร4 matrix โ (pair_state, giver_position) -----
+
+_MATRIX: dict[tuple[str, str], tuple[str, str | None]] = {
+ ("up", "up"): ("dancing", None),
+ ("up", "calm"): ("cheering", "A"),
+ ("up", "down"): ("comforting", "A"),
+ ("up", "tense"): ("defusing", "A"),
+
+ ("calm", "up"): ("cheering", "B"),
+ ("calm", "calm"): ("idle", None),
+ ("calm", "down"): ("comforting", "A"),
+ ("calm", "tense"): ("listening", "A"),
+
+ ("down", "up"): ("comforting", "B"),
+ ("down", "calm"): ("comforting", "B"),
+ ("down", "down"): ("sitting_together", None),
+ ("down", "tense"): ("tension", None),
+
+ ("tense", "up"): ("defusing", "B"),
+ ("tense", "calm"): ("listening", "B"),
+ ("tense", "down"): ("tension", None),
+ ("tense", "tense"): ("back_turned", None),
+}
+
+
+def resolve_pair(mood_a: str, mood_b: str) -> tuple[str, str | None]:
+ """Look up (pair_state, giver) for the given mood pair. Fallback: ('idle', None)."""
+ return _MATRIX.get((mood_a, mood_b), ("idle", None))
+
+
+SOLO_STATE_MAP = {
+ "joy": "joy",
+ "sadness": "sadness",
+ "anger": "anger",
+ "fear": "fear",
+ "surprise": "surprise",
+ "disgust": "disgust",
+ "neutral": "neutral",
+ "other": "neutral",
+ "unknown": "neutral",
+}
+
+
+def map_solo_state(emotion: str) -> str:
+ return SOLO_STATE_MAP.get(emotion.lower().strip(), "neutral")
+
+
+def select_representative_emotion(summary: SpeakerSummary, neutral_threshold: float = 0.70) -> str:
+ """Choose the speaker's representative emotion for display.
+
+ If neutral >= threshold, return neutral. Otherwise return the top non-neutral emotion.
+ Prevents conversations with mostly-neutral segments from always showing as neutral
+ when there are meaningful emotional moments.
+ """
+ dist = summary.emotion_distribution or {}
+ if not dist:
+ return summary.dominant_emotion or "neutral"
+
+ neutral_ratio = dist.get("neutral", 0.0)
+ if neutral_ratio >= neutral_threshold:
+ return "neutral"
+
+ non_neutral = {k: v for k, v in dist.items() if k != "neutral"}
+ if not non_neutral:
+ return "neutral"
+ return max(non_neutral, key=non_neutral.get)
+
+
+def map_characters(speaker_summaries: dict[str, SpeakerSummary]) -> list[CharacterReaction]:
+ """Build CharacterReactions using the 4x4 mood matrix.
+
+ Emits pair_state (9 symmetric labels) + role ("giver"/"receiver"/None).
+ solo_state remains the 7-class representative emotion for display.
+ """
+ speakers = sorted(speaker_summaries.keys())
+
+ if len(speakers) < 2:
+ sp = speakers[0] if speakers else "speaker_0"
+ summary = speaker_summaries.get(sp)
+ emotion = select_representative_emotion(summary) if summary else "neutral"
+ return [CharacterReaction(
+ speaker_id=sp,
+ solo_state=map_solo_state(emotion),
+ pair_state="idle",
+ role=None,
+ intensity=2,
+ )]
+
+ sp_a, sp_b = speakers[0], speakers[1]
+ sum_a = speaker_summaries[sp_a]
+ sum_b = speaker_summaries[sp_b]
+
+ em_a = select_representative_emotion(sum_a)
+ em_b = select_representative_emotion(sum_b)
+
+ mood_a = mood(sum_a)
+ mood_b = mood(sum_b)
+ pair_state, giver = resolve_pair(mood_a, mood_b)
+
+ if giver == "A":
+ role_a, role_b = "giver", "receiver"
+ elif giver == "B":
+ role_a, role_b = "receiver", "giver"
+ else:
+ role_a = role_b = None
+
+ return [
+ CharacterReaction(speaker_id=sp_a, solo_state=map_solo_state(em_a), pair_state=pair_state, role=role_a, intensity=2),
+ CharacterReaction(speaker_id=sp_b, solo_state=map_solo_state(em_b), pair_state=pair_state, role=role_b, intensity=2),
+ ]
diff --git a/src/stage3/garden_logic.py b/src/stage3/garden_logic.py
new file mode 100644
index 0000000000000000000000000000000000000000..93cfe877b41bb66d4cbacbd59d0c6ded8aa27182
--- /dev/null
+++ b/src/stage3/garden_logic.py
@@ -0,0 +1,93 @@
+"""
+์ ์ ์ฑ์ฅ ๋ก์ง (Stage 3)
+
+ํตํ ํ์ง(๊ธ์ ๊ฐ์ ๋น์จ) โ growth_delta ๊ณ์ฐ
+๋์ ๋ ๋ฒจ ๊ด๋ฆฌ (max_level: 5)
+๋ฌด๋ ๊ฒฐ์ (happy / neglected / recovering / conflict)
+"""
+
+from src.common.schemas import GardenUpdate, SpeakerSummary
+
+MAX_LEVEL = 5
+
+POSITIVE_EMOTIONS = {"joy", "surprise"}
+NEGATIVE_EMOTIONS = {"anger", "sadness", "fear", "disgust"}
+
+
+def _positive_ratio(speaker_summaries: dict[str, SpeakerSummary]) -> float:
+ """์ ์ฒด ํ์์ ๊ฐ์ ๋ถํฌ์์ ๊ธ์ ๊ฐ์ ๋น์จ ๊ณ์ฐ."""
+ total_positive = 0.0
+ total_weight = 0.0
+
+ for summary in speaker_summaries.values():
+ for emotion, weight in summary.emotion_distribution.items():
+ total_weight += weight
+ if emotion.lower() in POSITIVE_EMOTIONS:
+ total_positive += weight
+
+ if total_weight == 0:
+ return 0.0
+ return total_positive / total_weight
+
+
+def _negative_ratio(speaker_summaries: dict[str, SpeakerSummary]) -> float:
+ """์ ์ฒด ํ์์ ๊ฐ์ ๋ถํฌ์์ ๋ถ์ ๊ฐ์ ๋น์จ ๊ณ์ฐ."""
+ total_negative = 0.0
+ total_weight = 0.0
+
+ for summary in speaker_summaries.values():
+ for emotion, weight in summary.emotion_distribution.items():
+ total_weight += weight
+ if emotion.lower() in NEGATIVE_EMOTIONS:
+ total_negative += weight
+
+ if total_weight == 0:
+ return 0.0
+ return total_negative / total_weight
+
+
+def calculate_growth(speaker_summaries: dict[str, SpeakerSummary]) -> int:
+ """ํตํ ํ์ง ๊ธฐ๋ฐ ์ฑ์ฅ ์ ์ (0-3)."""
+ pos = _positive_ratio(speaker_summaries)
+ neg = _negative_ratio(speaker_summaries)
+
+ if pos >= 0.5:
+ return 3 # ๋งค์ฐ ๊ธ์ ์ ํตํ
+ if pos >= 0.3 and neg < 0.3:
+ return 2 # ๊ธ์ ์ ํตํ
+ if neg < 0.5:
+ return 1 # ๋ณดํต ํตํ
+ return 0 # ๋ถ์ ์ ํตํ โ ์ฑ์ฅ ์์
+
+
+def determine_mood(
+ speaker_summaries: dict[str, SpeakerSummary],
+ current_level: int,
+) -> str:
+ """์ ์ ๋ฌด๋ ๊ฒฐ์ ."""
+ neg = _negative_ratio(speaker_summaries)
+ pos = _positive_ratio(speaker_summaries)
+
+ if neg >= 0.5:
+ return "conflict"
+ if neg >= 0.3 and pos < 0.3:
+ return "recovering"
+ if current_level <= 1 and pos < 0.3:
+ return "neglected"
+ return "happy"
+
+
+def compute_garden_update(
+ speaker_summaries: dict[str, SpeakerSummary],
+ current_level: int = 1,
+) -> GardenUpdate:
+ """Stage 2 ์ถ๋ ฅ์ผ๋ก๋ถํฐ ์ ์ ์
๋ฐ์ดํธ ๊ณ์ฐ."""
+ growth = calculate_growth(speaker_summaries)
+ new_level = min(MAX_LEVEL, current_level + growth)
+ mood = determine_mood(speaker_summaries, new_level)
+
+ return GardenUpdate(
+ growth_delta=growth,
+ total_level=new_level,
+ mood=mood,
+ )
diff --git a/src/stage3/process.py b/src/stage3/process.py
new file mode 100644
index 0000000000000000000000000000000000000000..32c428169669060ea55fdff3ce53c0fbb6a04afc
--- /dev/null
+++ b/src/stage3/process.py
@@ -0,0 +1,52 @@
+"""
+Stage 3 ์ค์ผ์คํธ๋ ์ดํฐ โ ์บ๋ฆญํฐ ๋ฐ์ + ์ ์ + ๋ฆฌ์บก ์์ฑ
+
+์
๋ ฅ: Stage 2 ์ถ๋ ฅ (Stage2Output)
+์ถ๋ ฅ: Stage 3 ์ถ๋ ฅ (Stage3Output)
+"""
+
+from src.common.schemas import Stage2Output, Stage3Output, Segment
+from src.stage3.character_mapping import map_characters
+from src.stage3.garden_logic import compute_garden_update
+from src.stage3.recap_generator import generate_recap_llm, generate_recap_fallback
+
+
+def process(
+ stage2_output: Stage2Output,
+ segments: list[Segment] | None = None,
+ current_garden_level: int = 1,
+ use_llm: bool = True,
+) -> Stage3Output:
+ """
+ Stage 3 ๋ฉ์ธ ์ง์
์ .
+
+ Args:
+ stage2_output: Stage 2์ ๊ฐ์ ๋ถ์ ๊ฒฐ๊ณผ
+ segments: Stage 1์ ๋ฐํ ์ธ๊ทธ๋จผํธ (๋ฆฌ์บก ์์ฑ์ ์ฌ์ฉ, ์์ผ๋ฉด fallback)
+ current_garden_level: ํ์ฌ ์ ์ ๋ ๋ฒจ (DB์์ ๊ฐ์ ธ์ด)
+ use_llm: True๋ฉด Claude API ๋ฆฌ์บก, False๋ฉด ๊ท์น ๊ธฐ๋ฐ fallback
+ """
+ # 1. ์บ๋ฆญํฐ ๋ฐ์ ๋งคํ
+ character_reactions = map_characters(stage2_output.speaker_summaries)
+
+ # 2. ์ ์ ์ฑ์ฅ ๊ณ์ฐ
+ garden_update = compute_garden_update(
+ stage2_output.speaker_summaries,
+ current_level=current_garden_level,
+ )
+
+ # 3. ๋ฆฌ์บก ์์ฑ
+ if use_llm and segments:
+ recap_card = generate_recap_llm(segments, stage2_output.speaker_summaries)
+ else:
+ recap_card = generate_recap_fallback(stage2_output.speaker_summaries)
+
+ return Stage3Output(
+ call_id=stage2_output.call_id,
+ character_reactions=character_reactions,
+ garden_update=garden_update,
+ recap_card=recap_card,
+ # Surface Stage2 data at top level for frontend convenience
+ emotions=stage2_output.emotions,
+ speaker_summaries=stage2_output.speaker_summaries,
+ )
diff --git a/src/stage3/recap_generator.py b/src/stage3/recap_generator.py
new file mode 100644
index 0000000000000000000000000000000000000000..098d7cb65590328337cbab7760dfb6cea1ddc4f0
--- /dev/null
+++ b/src/stage3/recap_generator.py
@@ -0,0 +1,239 @@
+"""
+LLM ๋ฆฌ์บก ์์ฑ (Stage 3)
+
+ํ๋กฌํํธ ์ค๊ณ: ์ฃผํ (PO)
+๊ตฌํ: ์๊ท (App Engineer)
+
+Claude API๋ก ํตํ ์์ฝ + ํ์ด๋ผ์ดํธ ์์ฑ.
+API ํค๊ฐ ์์ผ๋ฉด ๊ท์น ๊ธฐ๋ฐ fallback ์ฌ์ฉ.
+"""
+
+import os
+
+try:
+ import anthropic
+except ImportError:
+ anthropic = None # type: ignore[assignment]
+
+from src.common.schemas import RecapCard, SpeakerSummary, Segment
+
+SYSTEM_PROMPT = """You are a garden keeper observing a couple's conversation for the UsTwo app.
+Your role is to notice what bloomed between them โ not to judge or diagnose.
+
+Output JSON with exactly these fields:
+
+- "title": A warm, specific title that makes THIS moment recognizable at a glance
+ in a list of recent calls (max 50 chars). Combine two ingredients:
+ (1) ONE concrete hook from the transcript โ a topic they discussed, a decision,
+ a plan, a worry, a shared joke, or a specific object/event mentioned.
+ (2) A light garden-voice framing โ a verb or noun like bloomed, warmth, roots,
+ quiet, seed, shade, first light, steady, shelter.
+ DO NOT use generic combo templates like "A warm moment together" or "A brave, honest
+ moment" โ every call ends up with the same title and the reader can't tell them
+ apart. The concrete hook is what makes titles distinct.
+
+ Examples (illustrative patterns, do not copy verbatim):
+ EN โ "Warmth bloomed over weekend plans"
+ "Quiet roots for a hard week"
+ "Light over the broken fridge"
+ "Steady ground when work got heavy"
+ "A seed of a small apology"
+ KO โ "์ฃผ๋ง ๊ณํ์ ํผ์ด๋ ๋ฐ๋ปํจ"
+ "ํ๋ ํ ์ฃผ์ ๋ด๋ฆฐ ์กฐ์ฉํ ๋ฟ๋ฆฌ"
+ "๊ณ ์ฅ ๋ ๋์ฅ๊ณ ์ ๋ฒ์ง ์์"
+ "๋ฐ์ ํ๋ฃจ ๊ณ์ ๋จธ๋ฌธ ์๊ฐ"
+ "์์ ๋ฏธ์ํจ์ด ์ฌ์ด์ง ์๊ฐ"
+
+- "summary": 2-3 sentences observing what happened. Use "you two", "together",
+ never "speakers" or "participants".
+
+- "highlights": 2-4 key moments using the phrase library:
+ Laughter โ "Laughter bloomed often"
+ Emotion match โ "Your hearts were in tune"
+ Good listening โ "One of you listened with great care"
+ Conflict resolved โ "Honesty brought you closer"
+ Long silence โ "Silence was shared, not empty"
+ Worry detected โ "Care for each other came through"
+
+Rules:
+- Title must reference something specific from THIS conversation, not a template.
+- Observe, don't judge. All emotions are valid.
+- Slight warmth bias โ frame everything as growth, never as failure.
+- Never use: "negative", "tense", "sad", "anxious" as direct labels.
+- Never use: "speakers", "participants", "users", "energy", "vibes", "detected".
+- Use garden language: bloom, grow, warmth, sunshine, roots, together.
+- Write in the language matching the majority of the transcript (Korean or English).
+- For Korean, use: ์๊ฐ/์๊ฐ (not ํตํ), ๋/๋ ์ฌ๋ (not speakers), ํผ์ด๋๋ค/์๋ผ๋ค (garden verbs).
+"""
+
+
+def _build_transcript_text(segments: list[Segment]) -> str:
+ """์ธ๊ทธ๋จผํธ ๋ฆฌ์คํธ๋ฅผ LLM ์
๋ ฅ์ฉ ํ
์คํธ๋ก ๋ณํ."""
+ lines = []
+ for seg in segments:
+ lines.append(f"[{seg.speaker_id}] ({seg.start:.1f}s-{seg.end:.1f}s): {seg.text}")
+ return "\n".join(lines)
+
+
+def _build_emotion_summary(speaker_summaries: dict[str, SpeakerSummary]) -> str:
+ """๊ฐ์ ์์ฝ ํ
์คํธ ์์ฑ."""
+ parts = []
+ for sp_id, summary in sorted(speaker_summaries.items()):
+ top_emotions = sorted(summary.emotion_distribution.items(), key=lambda x: -x[1])[:3]
+ emotions_str = ", ".join(f"{e}({v:.0%})" for e, v in top_emotions)
+ parts.append(f"{sp_id}: dominant={summary.dominant_emotion}, distribution=[{emotions_str}]")
+ return "\n".join(parts)
+
+
+def generate_recap_llm(
+ segments: list[Segment],
+ speaker_summaries: dict[str, SpeakerSummary],
+) -> RecapCard:
+ """Claude API๋ก ๋ฆฌ์บก ์์ฑ. API ํค๊ฐ ์์ผ๋ฉด fallback."""
+ api_key = os.environ.get("ANTHROPIC_API_KEY")
+ if not api_key:
+ return generate_recap_fallback(speaker_summaries)
+
+ try:
+ import anthropic
+ import json
+
+ client = anthropic.Anthropic(api_key=api_key)
+
+ transcript = _build_transcript_text(segments)
+ emotions = _build_emotion_summary(speaker_summaries)
+
+ user_message = f"""Transcript:
+{transcript}
+
+Emotion analysis:
+{emotions}
+
+Generate the recap card as JSON."""
+
+ response = client.messages.create(
+ model="claude-sonnet-4-20250514",
+ max_tokens=500,
+ system=SYSTEM_PROMPT,
+ messages=[{"role": "user", "content": user_message}],
+ )
+
+ text = response.content[0].text
+ # Extract JSON from response (handle markdown code blocks)
+ if "```" in text:
+ text = text.split("```")[1]
+ if text.startswith("json"):
+ text = text[4:]
+ data = json.loads(text.strip())
+
+ return RecapCard(
+ title=data.get("title", "Call Recap"),
+ summary=data.get("summary", ""),
+ highlights=data.get("highlights", []),
+ )
+ except Exception:
+ return generate_recap_fallback(speaker_summaries)
+
+
+def generate_recap_fallback(speaker_summaries: dict[str, SpeakerSummary]) -> RecapCard:
+ """Garden-worldview fallback recap (when API key is absent)."""
+ speakers = sorted(speaker_summaries.keys())
+ if not speakers:
+ return RecapCard(title="A Quiet Moment", summary="A moment passed between you two.", highlights=[])
+
+ dominant_emotions = [speaker_summaries[sp].dominant_emotion for sp in speakers]
+
+ positive = {"joy", "surprise"}
+ negative = {"anger", "sadness", "fear", "disgust"}
+
+ pos_count = sum(1 for e in dominant_emotions if e in positive)
+ neg_count = sum(1 for e in dominant_emotions if e in negative)
+
+ if pos_count >= len(speakers):
+ title = "A moment that felt like sunshine"
+ summary = "Warmth filled this moment. You two were in tune, and it showed."
+ highlights = ["Laughter bloomed often", "Your hearts were in tune"]
+ elif neg_count >= len(speakers):
+ title = "A brave, honest moment"
+ summary = "This took courage. You two showed up for each other, even when it was hard."
+ highlights = ["Honesty brought you closer", "Care for each other came through"]
+ elif pos_count > 0 and neg_count > 0:
+ title = "One heart held the other"
+ summary = "One of you carried warmth while the other worked through something. That balance matters."
+ highlights = ["One of you listened with great care", "Care for each other came through"]
+ else:
+ title = "Just being together was enough"
+ summary = "A calm, easy moment. Sometimes the best conversations are the quiet ones."
+ highlights = ["Silence was shared, not empty", "Your hearts were in tune"]
+
+ labels = {speakers[0]: "You"} if len(speakers) == 1 else {speakers[0]: "You", speakers[1]: "Your partner"}
+ for sp in speakers:
+ label = labels.get(sp, sp)
+ em = speaker_summaries[sp].dominant_emotion
+ if em in positive:
+ highlights.append(f"{label} brought warmth to this moment")
+ elif em in negative:
+ highlights.append(f"{label} {'were' if label == 'You' else 'was'} working through something real")
+ else:
+ highlights.append(f"{label} {'were' if label == 'You' else 'was'} steady and present")
+
+ return RecapCard(title=title, summary=summary, highlights=highlights[:4])
+
+
+def generate_recap_from_transcript(transcript: str) -> RecapCard:
+ """Generate recap from raw transcript only (no emotion data).
+
+ Used for quick preview โ Whisper API transcript before full pipeline runs.
+ Falls back to a generic recap if no API key or on error.
+ """
+ if not transcript.strip():
+ return RecapCard(
+ title="A quiet moment together",
+ summary="A moment passed between you two.",
+ highlights=[],
+ )
+
+ api_key = os.environ.get("ANTHROPIC_API_KEY")
+ if not api_key:
+ return RecapCard(
+ title="A moment between you two",
+ summary="Your conversation is being analyzed. The full picture will bloom soon.",
+ highlights=["A seed has been planted"],
+ )
+
+ try:
+ import json
+
+ client = anthropic.Anthropic(api_key=api_key)
+
+ user_message = f"""Transcript:
+{transcript}
+
+Generate the recap card as JSON. Note: emotion analysis is not yet available โ
+base your recap only on the conversation content and tone."""
+
+ response = client.messages.create(
+ model="claude-sonnet-4-20250514",
+ max_tokens=500,
+ system=SYSTEM_PROMPT,
+ messages=[{"role": "user", "content": user_message}],
+ )
+
+ text = response.content[0].text
+ if "```" in text:
+ text = text.split("```")[1]
+ if text.startswith("json"):
+ text = text[4:]
+ data = json.loads(text.strip())
+
+ return RecapCard(
+ title=data.get("title", "A moment together"),
+ summary=data.get("summary", ""),
+ highlights=data.get("highlights", []),
+ )
+ except Exception:
+ return RecapCard(
+ title="A moment between you two",
+ summary="Your conversation is being analyzed. The full picture will bloom soon.",
+ highlights=["A seed has been planted"],
+ )
diff --git a/src/stage3/requirements.txt b/src/stage3/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..bd989ed78a10c520fc716062bce80e42ab6743d3
--- /dev/null
+++ b/src/stage3/requirements.txt
@@ -0,0 +1,3 @@
+anthropic>=0.25.0
+openai>=1.12.0
+pydantic>=2.0.0
diff --git a/src/stage4/__init__.py b/src/stage4/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/src/stage4/database.py b/src/stage4/database.py
new file mode 100644
index 0000000000000000000000000000000000000000..3e3e0c969b158529f1ba6b3056054b9be5ee3eae
--- /dev/null
+++ b/src/stage4/database.py
@@ -0,0 +1,36 @@
+"""
+UsTwo SQLite database setup.
+
+Uses SQLAlchemy with SQLite โ no external DB server needed.
+DB file: data/ustwo.db (created automatically on first run).
+"""
+import os
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker, DeclarativeBase
+
+DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "data")
+DATABASE_URL = f"sqlite:///{os.path.join(DATA_DIR, 'ustwo.db')}"
+
+
+class Base(DeclarativeBase):
+ pass
+
+
+engine = create_engine(DATABASE_URL, echo=False)
+SessionLocal = sessionmaker(bind=engine)
+
+
+def init_db():
+ """Create all tables. Safe to call multiple times."""
+ os.makedirs(DATA_DIR, exist_ok=True)
+ from . import models # noqa: F401 โ ensure models are registered
+ Base.metadata.create_all(bind=engine)
+
+
+def get_db():
+ """FastAPI dependency โ yields a DB session, closes after request."""
+ db = SessionLocal()
+ try:
+ yield db
+ finally:
+ db.close()
diff --git a/src/stage4/main.py b/src/stage4/main.py
new file mode 100644
index 0000000000000000000000000000000000000000..fa3f21c99fd75f9e05783550160cd7b4995a817d
--- /dev/null
+++ b/src/stage4/main.py
@@ -0,0 +1,682 @@
+import json
+import logging
+import random
+import threading
+import uuid
+from datetime import datetime, timezone
+from pathlib import Path
+
+from fastapi import FastAPI, File, UploadFile, HTTPException, Depends
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from sqlalchemy.orm import Session
+
+from .database import init_db, get_db, SessionLocal
+from . import models
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s")
+logger = logging.getLogger(__name__)
+
+app = FastAPI(title="UsTwo API", version="0.3.0")
+
+app.add_middleware(
+ CORSMiddleware,
+ allow_origins=["*"],
+ allow_credentials=True,
+ allow_methods=["*"],
+ allow_headers=["*"],
+)
+
+UPLOAD_DIR = Path("data/samples")
+UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
+
+ALLOWED_EXTENSIONS = {".wav", ".m4a", ".mp3", ".ogg"}
+MAX_UPLOAD_BYTES = 50 * 1024 * 1024 # 50 MB
+
+
+@app.on_event("startup")
+def on_startup():
+ init_db()
+ _seed_demo_calls()
+
+
+def _seed_demo_calls():
+ """Pre-create call records for demo test scenarios (data/samples/*.json)."""
+ db = next(get_db())
+ demo_ids = ["test_happy", "test_tense", "test_mixed", "test_neutral"]
+ for call_id in demo_ids:
+ existing = db.query(models.Call).filter_by(id=call_id).first()
+ if not existing:
+ db.add(models.Call(id=call_id, audio_path="", status="uploaded"))
+ db.commit()
+
+
+# โโโ Health โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+@app.get("/api/health")
+def health():
+ return {"status": "ok", "timestamp": datetime.now(timezone.utc).isoformat()}
+
+
+# โโโ Upload โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+@app.post("/api/upload")
+async def upload_audio(file: UploadFile = File(...), db: Session = Depends(get_db)):
+ if not file.filename:
+ raise HTTPException(status_code=400, detail="No filename provided")
+
+ ext = Path(file.filename).suffix.lower()
+ if ext not in ALLOWED_EXTENSIONS:
+ raise HTTPException(
+ status_code=400,
+ detail=f"Unsupported file type: {ext}. Allowed: {ALLOWED_EXTENSIONS}",
+ )
+
+ content = await file.read()
+ if len(content) == 0:
+ raise HTTPException(status_code=400, detail="Empty file")
+ if len(content) > MAX_UPLOAD_BYTES:
+ raise HTTPException(
+ status_code=400,
+ detail=f"File too large ({len(content) // (1024*1024)}MB). Max: {MAX_UPLOAD_BYTES // (1024*1024)}MB",
+ )
+
+ call_id = f"call_{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:6]}"
+ save_path = UPLOAD_DIR / f"{call_id}{ext}"
+ save_path.write_bytes(content)
+
+ # Save to DB
+ db.add(models.Call(id=call_id, audio_path=str(save_path)))
+ db.commit()
+
+ return {"status": "success", "call_id": call_id, "filename": save_path.name}
+
+
+# โโโ Pipeline runner โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def _run_full_pipeline(audio_path: str, call_id: str):
+ """Run Stage 1 โ 2 โ 3 on a real audio file.
+
+ Returns (stage3_output, stage1_output, stage2_output) tuple.
+ Raises RuntimeError if ML dependencies are missing.
+ """
+ from src.stage1.process import process as stage1_process
+ from src.stage2.process import process as stage2_process
+ from src.stage3.process import process as stage3_process
+
+ logger.info("Pipeline start: %s (%s)", call_id, audio_path)
+
+ # Stage 1: Diarization + ASR
+ stage1_out = stage1_process(audio_path)
+ logger.info(
+ "Stage 1 done: %d segments, %.1fs",
+ len(stage1_out.segments),
+ stage1_out.processing_info.processing_time_sec,
+ )
+
+ # Stage 2: Emotion analysis
+ stage2_out = stage2_process(stage1_out)
+ logger.info(
+ "Stage 2 done: %d emotions, speakers=%s",
+ len(stage2_out.emotions),
+ list(stage2_out.speaker_summaries.keys()),
+ )
+
+ # Stage 3: Character + Garden + Recap
+ use_llm = False # rule-based fallback for now
+ stage3_out = stage3_process(
+ stage2_out,
+ segments=stage1_out.segments,
+ use_llm=use_llm,
+ )
+ logger.info("Stage 3 done: %s", call_id)
+
+ return stage3_out, stage1_out, stage2_out
+
+
+_MOCK_SCENARIOS = [
+ # (sp0_emotion, sp1_emotion) โ cycles through for consistent demo results
+ ("joy", "joy"),
+ ("joy", "sadness"),
+ ("anger", "sadness"),
+ ("neutral", "neutral"),
+ ("surprise", "joy"),
+]
+_mock_index = 0
+
+
+def _generate_mock_stage2(call_id: str):
+ """Mock Stage 2 output โ deterministic scenarios with rich segment data for Landscape."""
+ global _mock_index
+ from src.common.schemas import Stage2Output, EmotionResult, SpeakerSummary
+
+ sp0_emotion, sp1_emotion = _MOCK_SCENARIOS[_mock_index % len(_MOCK_SCENARIOS)]
+ _mock_index += 1
+
+ def _fixed_summary(emotion: str) -> SpeakerSummary:
+ return SpeakerSummary(
+ dominant_emotion=emotion,
+ emotion_distribution={emotion: 0.70, "neutral": 0.30},
+ avg_confidence=0.82,
+ )
+
+ sp0_summary = _fixed_summary(sp0_emotion)
+ sp1_summary = _fixed_summary(sp1_emotion)
+
+ # Generate 6-8 alternating segments for a visually rich Emotional Landscape
+ _SEGMENT_PATTERNS = {
+ ("joy", "joy"): [
+ ("speaker_0", "joy", 0.85), ("speaker_1", "joy", 0.78),
+ ("speaker_0", "surprise", 0.65), ("speaker_1", "joy", 0.82),
+ ("speaker_0", "joy", 0.88), ("speaker_1", "neutral", 0.70),
+ ("speaker_0", "joy", 0.80), ("speaker_1", "joy", 0.75),
+ ],
+ ("joy", "sadness"): [
+ ("speaker_0", "joy", 0.80), ("speaker_1", "neutral", 0.72),
+ ("speaker_0", "joy", 0.75), ("speaker_1", "sadness", 0.68),
+ ("speaker_0", "neutral", 0.70), ("speaker_1", "sadness", 0.78),
+ ("speaker_0", "joy", 0.82), ("speaker_1", "sadness", 0.65),
+ ],
+ ("anger", "sadness"): [
+ ("speaker_0", "neutral", 0.72), ("speaker_1", "neutral", 0.70),
+ ("speaker_0", "anger", 0.78), ("speaker_1", "sadness", 0.74),
+ ("speaker_0", "anger", 0.82), ("speaker_1", "fear", 0.65),
+ ("speaker_0", "neutral", 0.68), ("speaker_1", "sadness", 0.80),
+ ("speaker_0", "anger", 0.75), ("speaker_1", "sadness", 0.72),
+ ],
+ ("neutral", "neutral"): [
+ ("speaker_0", "neutral", 0.85), ("speaker_1", "neutral", 0.82),
+ ("speaker_0", "neutral", 0.78), ("speaker_1", "joy", 0.60),
+ ("speaker_0", "neutral", 0.80), ("speaker_1", "neutral", 0.75),
+ ("speaker_0", "joy", 0.62), ("speaker_1", "neutral", 0.80),
+ ],
+ ("surprise", "joy"): [
+ ("speaker_0", "neutral", 0.72), ("speaker_1", "joy", 0.75),
+ ("speaker_0", "surprise", 0.80), ("speaker_1", "joy", 0.82),
+ ("speaker_0", "surprise", 0.85), ("speaker_1", "surprise", 0.70),
+ ("speaker_0", "joy", 0.78), ("speaker_1", "joy", 0.80),
+ ],
+ }
+ pattern = _SEGMENT_PATTERNS.get((sp0_emotion, sp1_emotion), _SEGMENT_PATTERNS[("neutral", "neutral")])
+
+ emotions = []
+ for i, (spk, emo, conf) in enumerate(pattern):
+ emotions.append(EmotionResult(
+ speaker_id=spk, segment_id=i,
+ audio_emotion=emo, audio_confidence=conf,
+ text_emotion=emo, text_confidence=max(0.5, conf - 0.1),
+ fused_emotion=emo, fused_confidence=conf,
+ ))
+
+ return Stage2Output(
+ call_id=call_id,
+ emotions=emotions,
+ speaker_summaries={
+ "speaker_0": sp0_summary,
+ "speaker_1": sp1_summary,
+ },
+ )
+
+
+# โโโ Background pipeline worker โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+def _quick_recap(call_id: str, audio_path: str, db: Session):
+ """Phase 1: Fast recap via Whisper API + Claude (~15s).
+
+ Stores partial result with status 'preview'. Requires OPENAI_API_KEY.
+ """
+ import os
+ openai_key = os.environ.get("OPENAI_API_KEY")
+ if not openai_key:
+ logger.info("OPENAI_API_KEY not set, skipping quick recap for %s", call_id)
+ return
+
+ from openai import OpenAI
+ from src.stage3.recap_generator import generate_recap_from_transcript
+
+ logger.info("Quick recap start: %s", call_id)
+
+ # 1. Whisper API transcription
+ openai_client = OpenAI(api_key=openai_key)
+ with open(audio_path, "rb") as f:
+ transcription = openai_client.audio.transcriptions.create(
+ model="whisper-1",
+ file=f,
+ )
+ transcript = transcription.text
+ logger.info("Whisper API done: %d chars", len(transcript))
+
+ # 2. Claude recap from transcript
+ recap_card = generate_recap_from_transcript(transcript)
+
+ # 3. Build partial result with default character reactions
+ partial_result = {
+ "call_id": call_id,
+ "character_reactions": [
+ {"speaker_id": "speaker_0", "solo_state": "neutral", "pair_state": "sitting_together"},
+ {"speaker_id": "speaker_1", "solo_state": "neutral", "pair_state": "sitting_together"},
+ ],
+ "garden_update": {"growth_delta": 0, "total_level": 1, "mood": "happy"},
+ "recap_card": recap_card.model_dump(),
+ }
+
+ # 4. Store in DB
+ db.add(models.AnalysisResult(
+ call_id=call_id,
+ stage3_json=json.dumps(partial_result),
+ blue_emotion="neutral",
+ pink_emotion="neutral",
+ garden_delta=0,
+ ))
+ call = db.query(models.Call).filter(models.Call.id == call_id).first()
+ if call:
+ call.status = "preview"
+ db.commit()
+
+ logger.info("Quick recap done: %s โ '%s'", call_id, recap_card.title)
+
+
+def _run_pipeline_background(call_id: str, audio_path: str):
+ """Run 2-phase pipeline in a background thread.
+
+ Phase 1: Quick recap via Whisper API + Claude (~15s) โ status 'preview'
+ Phase 2: Full ML pipeline (diarization + emotion) โ status 'done'
+ Phase 2 preserves Phase 1 recap, only updates character reactions + garden.
+ """
+ db = SessionLocal()
+ try:
+ call = db.query(models.Call).filter(models.Call.id == call_id).first()
+ if not call:
+ return
+
+ call.status = "analyzing"
+ db.commit()
+
+ # Phase 1: Quick Recap (best-effort, failure doesn't block Phase 2)
+ has_preview = False
+ try:
+ _quick_recap(call_id, audio_path, db)
+ has_preview = True
+ except Exception as e:
+ logger.warning("Quick recap failed for %s: %s", call_id, e)
+
+ # Phase 2: Full ML Pipeline
+ try:
+ import sys
+ print(f"[PIPELINE] Starting full pipeline for {call_id}", file=sys.stderr, flush=True)
+ stage3_result, stage1_out, stage2_out = _run_full_pipeline(audio_path, call_id)
+ pipeline_mode = "full"
+ call = db.query(models.Call).filter(models.Call.id == call_id).first()
+ call.duration = stage1_out.duration
+ except ImportError as e:
+ logger.warning("ML deps missing (%s), falling back to mock", e)
+ from src.stage3.process import process as stage3_process
+ stage2_out = _generate_mock_stage2(call_id)
+ stage3_result = stage3_process(stage2_out, use_llm=False)
+ pipeline_mode = "mock"
+ call = db.query(models.Call).filter(models.Call.id == call_id).first()
+
+ result_dict = stage3_result.model_dump()
+ result_dict["emotions"] = [e.model_dump() for e in stage2_out.emotions]
+ result_dict["stage2_output"] = stage2_out.model_dump()
+
+ from src.stage3.character_mapping import select_representative_emotion
+ summaries = stage2_out.speaker_summaries or {}
+ blue_emo = select_representative_emotion(summaries["speaker_0"]) if "speaker_0" in summaries else None
+ pink_emo = select_representative_emotion(summaries["speaker_1"]) if "speaker_1" in summaries else None
+
+ if has_preview:
+ # Preserve Phase 1 recap, update only character reactions + garden
+ existing = db.query(models.AnalysisResult).filter(
+ models.AnalysisResult.call_id == call_id
+ ).first()
+ if existing:
+ preview_data = json.loads(existing.stage3_json)
+ result_dict["recap_card"] = preview_data["recap_card"]
+ existing.stage3_json = json.dumps(result_dict)
+ existing.blue_emotion = blue_emo
+ existing.pink_emotion = pink_emo
+ existing.garden_delta = result_dict.get("garden_update", {}).get("growth_delta", 0)
+ else:
+ db.add(models.AnalysisResult(
+ call_id=call_id,
+ stage3_json=json.dumps(result_dict),
+ blue_emotion=blue_emo,
+ pink_emotion=pink_emo,
+ garden_delta=result_dict.get("garden_update", {}).get("growth_delta", 0),
+ ))
+ else:
+ db.add(models.AnalysisResult(
+ call_id=call_id,
+ stage3_json=json.dumps(result_dict),
+ blue_emotion=blue_emo,
+ pink_emotion=pink_emo,
+ garden_delta=result_dict.get("garden_update", {}).get("growth_delta", 0),
+ ))
+
+ _update_garden(db, result_dict)
+ call.status = "done"
+ db.commit()
+
+ logger.info("Background pipeline done: %s (mode=%s, preview=%s)", call_id, pipeline_mode, has_preview)
+
+ except Exception as e:
+ logger.error("Background pipeline failed for %s: %s", call_id, e, exc_info=True)
+ db.rollback()
+ try:
+ call = db.query(models.Call).filter(models.Call.id == call_id).first()
+ if call:
+ call.status = "error"
+ call.error_message = str(e)
+ db.commit()
+ except Exception:
+ db.rollback()
+ finally:
+ db.close()
+
+
+# โโโ Analyze โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+@app.post("/api/analyze", status_code=202)
+def analyze(call_id: str, db: Session = Depends(get_db)):
+ """Start async analysis pipeline. Returns immediately with 202.
+
+ Poll GET /api/analyze/{call_id}/status for progress.
+ When status is 'done', result is available at GET /api/calls/{call_id}.
+ """
+ call = db.query(models.Call).filter(models.Call.id == call_id).first()
+ if not call:
+ raise HTTPException(status_code=404, detail="Call not found")
+
+ if call.status in ("analyzing", "preview"):
+ return {"status": "analyzing", "call_id": call_id, "message": "Already in progress"}
+
+ if call.status == "done":
+ # Already analyzed โ return existing result
+ existing = db.query(models.AnalysisResult).filter(
+ models.AnalysisResult.call_id == call_id
+ ).first()
+ if existing:
+ return {
+ "status": "done",
+ "call_id": call_id,
+ "result": json.loads(existing.stage3_json),
+ }
+
+ audio_path = call.audio_path
+
+ # Check for pre-computed Stage 2 JSON (์น์ฌ's pipeline output or test data)
+ stage2_json_path = Path("data") / f"{call_id}_stage2.json"
+ if not stage2_json_path.exists():
+ stage2_json_path = Path("data/samples") / f"{call_id}_stage2.json"
+
+ if stage2_json_path.exists():
+ # Use real Stage 2 output โ run Stage 3 only
+ from src.common.schemas import Stage2Output
+ from src.stage3.process import process as stage3_process
+ try:
+ stage2 = Stage2Output.model_validate_json(stage2_json_path.read_text())
+ except Exception as e:
+ raise HTTPException(status_code=400, detail=f"Invalid Stage 2 JSON: {e}")
+ import os
+ has_api_key = bool(os.environ.get("ANTHROPIC_API_KEY"))
+ stage3_result = stage3_process(stage2, use_llm=has_api_key)
+ result_dict = stage3_result.model_dump()
+ result_dict["emotions"] = [e.model_dump() for e in stage2.emotions]
+
+ from src.stage3.character_mapping import select_representative_emotion
+ summaries = stage2.speaker_summaries or {}
+ blue_emo = select_representative_emotion(summaries["speaker_0"]) if "speaker_0" in summaries else None
+ pink_emo = select_representative_emotion(summaries["speaker_1"]) if "speaker_1" in summaries else None
+ db.add(models.AnalysisResult(
+ call_id=call_id,
+ stage3_json=json.dumps(result_dict),
+ blue_emotion=blue_emo,
+ pink_emotion=pink_emo,
+ garden_delta=result_dict.get("garden_update", {}).get("growth_delta", 0),
+ ))
+ _update_garden(db, result_dict)
+ call.status = "done"
+ db.commit()
+
+ return {
+ "status": "done",
+ "call_id": call_id,
+ "pipeline_mode": "stage2_json",
+ "result": result_dict,
+ }
+
+ if not audio_path or not Path(audio_path).exists():
+ # No audio file and no stage2 JSON โ run mock synchronously
+ from src.stage3.process import process as stage3_process
+ stage2 = _generate_mock_stage2(call_id)
+ import os
+ has_api_key = bool(os.environ.get("ANTHROPIC_API_KEY"))
+ stage3_result = stage3_process(stage2, use_llm=has_api_key)
+ result_dict = stage3_result.model_dump()
+ result_dict["emotions"] = [e.model_dump() for e in stage2.emotions]
+
+ from src.stage3.character_mapping import select_representative_emotion
+ summaries = stage2.speaker_summaries or {}
+ blue_emo = select_representative_emotion(summaries["speaker_0"]) if "speaker_0" in summaries else None
+ pink_emo = select_representative_emotion(summaries["speaker_1"]) if "speaker_1" in summaries else None
+ db.add(models.AnalysisResult(
+ call_id=call_id,
+ stage3_json=json.dumps(result_dict),
+ blue_emotion=blue_emo,
+ pink_emotion=pink_emo,
+ garden_delta=result_dict.get("garden_update", {}).get("growth_delta", 0),
+ ))
+ _update_garden(db, result_dict)
+ call.status = "done"
+ db.commit()
+
+ return {
+ "status": "done",
+ "call_id": call_id,
+ "pipeline_mode": "mock",
+ "result": result_dict,
+ }
+
+ # Launch background pipeline
+ thread = threading.Thread(
+ target=_run_pipeline_background,
+ args=(call_id, audio_path),
+ daemon=True,
+ )
+ thread.start()
+
+ return {"status": "analyzing", "call_id": call_id, "message": "Pipeline started"}
+
+
+@app.get("/api/analyze/{call_id}/status")
+def analyze_status(call_id: str, db: Session = Depends(get_db)):
+ """Poll analysis progress."""
+ call = db.query(models.Call).filter(models.Call.id == call_id).first()
+ if not call:
+ raise HTTPException(status_code=404, detail="Call not found")
+
+ if call.status in ("done", "preview"):
+ result = db.query(models.AnalysisResult).filter(
+ models.AnalysisResult.call_id == call_id
+ ).first()
+ return {
+ "status": call.status,
+ "call_id": call_id,
+ "result": json.loads(result.stage3_json) if result else None,
+ }
+
+ if call.status == "error":
+ return {
+ "status": "error",
+ "call_id": call_id,
+ "error": call.error_message,
+ }
+
+ return {"status": call.status, "call_id": call_id}
+
+
+# โโโ Calls (history) โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+@app.get("/api/calls")
+def list_calls(db: Session = Depends(get_db)):
+ results = (
+ db.query(models.AnalysisResult, models.Call.status)
+ .join(models.Call, models.AnalysisResult.call_id == models.Call.id)
+ .order_by(models.AnalysisResult.created_at.desc())
+ .limit(50)
+ .all()
+ )
+
+ def _extract_title(stage3_json: str | None) -> str | None:
+ if not stage3_json:
+ return None
+ try:
+ parsed = json.loads(stage3_json)
+ except (ValueError, TypeError):
+ return None
+ recap = parsed.get("recap_card") if isinstance(parsed, dict) else None
+ if not isinstance(recap, dict):
+ return None
+ title = recap.get("title")
+ return title if isinstance(title, str) and title.strip() else None
+
+ return [
+ {
+ "call_id": r.AnalysisResult.call_id,
+ "blue_emotion": r.AnalysisResult.blue_emotion if r.status == "done" else None,
+ "pink_emotion": r.AnalysisResult.pink_emotion if r.status == "done" else None,
+ "garden_delta": r.AnalysisResult.garden_delta,
+ "created_at": r.AnalysisResult.created_at.isoformat() if r.AnalysisResult.created_at else None,
+ "status": r.status,
+ "recap_title": _extract_title(r.AnalysisResult.stage3_json) if r.status == "done" else None,
+ }
+ for r in results
+ ]
+
+
+@app.get("/api/calls/{call_id}")
+def get_call(call_id: str, db: Session = Depends(get_db)):
+ result = db.query(models.AnalysisResult).filter(models.AnalysisResult.call_id == call_id).first()
+ if not result:
+ raise HTTPException(status_code=404, detail="Call not found")
+ return {
+ "call_id": result.call_id,
+ "result": json.loads(result.stage3_json),
+ "created_at": result.created_at.isoformat() if result.created_at else None,
+ }
+
+
+# โโโ Check-ins โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+VALID_LEVELS = {"deep", "warm", "growing", "different", "learning"}
+
+
+class CheckInCreate(BaseModel):
+ iso_date: str
+ score: int
+ level: str
+ my_mood: str
+ partner_guess: str
+
+ def validate_fields(self):
+ if not (0 <= self.score <= 100):
+ raise HTTPException(status_code=400, detail=f"Score must be 0-100, got {self.score}")
+ if self.level not in VALID_LEVELS:
+ raise HTTPException(status_code=400, detail=f"Invalid level: {self.level}")
+
+
+@app.post("/api/checkins")
+def create_checkin(data: CheckInCreate, db: Session = Depends(get_db)):
+ data.validate_fields()
+ checkin = models.CheckIn(
+ iso_date=data.iso_date,
+ score=data.score,
+ level=data.level,
+ my_mood=data.my_mood,
+ partner_guess=data.partner_guess,
+ )
+ db.add(checkin)
+
+ # Update garden
+ _update_garden_from_checkin(db, data.score)
+ db.commit()
+ db.refresh(checkin)
+
+ return {"status": "success", "id": checkin.id}
+
+
+@app.get("/api/checkins")
+def list_checkins(db: Session = Depends(get_db)):
+ records = db.query(models.CheckIn).order_by(models.CheckIn.created_at.desc()).limit(100).all()
+ return [
+ {
+ "id": r.id,
+ "iso_date": r.iso_date,
+ "score": r.score,
+ "level": r.level,
+ "my_mood": r.my_mood,
+ "partner_guess": r.partner_guess,
+ "created_at": r.created_at.isoformat() if r.created_at else None,
+ }
+ for r in records
+ ]
+
+
+# โโโ Garden โโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโโ
+
+# design-system.md ยง4.6 thresholds
+def _compute_level(count: int) -> int:
+ if count >= 25: return 5
+ if count >= 15: return 4
+ if count >= 8: return 3
+ if count >= 3: return 2
+ return 1
+
+
+def _get_or_create_garden(db: Session) -> models.GardenState:
+ garden = db.query(models.GardenState).filter(models.GardenState.id == 1).first()
+ if not garden:
+ garden = models.GardenState(id=1)
+ db.add(garden)
+ db.flush()
+ return garden
+
+
+def _update_garden(db: Session, result_dict: dict):
+ garden = _get_or_create_garden(db)
+ garden.interaction_count += 1
+ garden.total_level = _compute_level(garden.interaction_count)
+ mood = result_dict.get("garden_update", {}).get("mood", "happy")
+ garden.last_mood = mood
+
+
+def _update_garden_from_checkin(db: Session, score: int):
+ garden = _get_or_create_garden(db)
+ garden.interaction_count += 1
+ garden.total_level = _compute_level(garden.interaction_count)
+ garden.last_mood = "happy" if score >= 60 else "recovering"
+
+
+@app.get("/api/garden")
+def get_garden(db: Session = Depends(get_db)):
+ garden = _get_or_create_garden(db)
+ return {
+ "interaction_count": garden.interaction_count,
+ "total_level": garden.total_level,
+ "last_mood": garden.last_mood,
+ }
+
+
+@app.put("/api/garden/interact")
+def garden_interact(positive: bool = True, db: Session = Depends(get_db)):
+ garden = _get_or_create_garden(db)
+ garden.interaction_count += 1
+ garden.total_level = _compute_level(garden.interaction_count)
+ garden.last_mood = "happy" if positive else "recovering"
+ db.commit()
+ return {
+ "interaction_count": garden.interaction_count,
+ "total_level": garden.total_level,
+ "last_mood": garden.last_mood,
+ }
diff --git a/src/stage4/models.py b/src/stage4/models.py
new file mode 100644
index 0000000000000000000000000000000000000000..2957031d95b3d67264c018738bdf03d504bac4bd
--- /dev/null
+++ b/src/stage4/models.py
@@ -0,0 +1,61 @@
+"""
+SQLAlchemy models for UsTwo.
+
+Tables:
+ calls โ uploaded audio files + metadata
+ analysis_results โ Stage 3 output per call
+ checkins โ empathy check-in records
+ garden_state โ singleton garden progression
+"""
+from datetime import datetime, timezone
+from sqlalchemy import Column, String, Integer, Float, Text, DateTime, ForeignKey
+from .database import Base
+
+
+def _utcnow() -> datetime:
+ return datetime.now(timezone.utc)
+
+
+class Call(Base):
+ __tablename__ = "calls"
+
+ id = Column(String, primary_key=True) # call_id (timestamp-based)
+ audio_path = Column(String, nullable=False)
+ duration = Column(Float, nullable=True)
+ status = Column(String, default="uploaded") # uploaded โ analyzing โ done โ error
+ error_message = Column(Text, nullable=True)
+ created_at = Column(DateTime, default=_utcnow)
+
+
+class AnalysisResult(Base):
+ __tablename__ = "analysis_results"
+
+ id = Column(Integer, primary_key=True, autoincrement=True)
+ call_id = Column(String, ForeignKey("calls.id"), nullable=False)
+ stage3_json = Column(Text, nullable=False) # full Stage3Output JSON
+ blue_emotion = Column(String, nullable=True) # quick-query columns
+ pink_emotion = Column(String, nullable=True)
+ garden_delta = Column(Integer, default=0)
+ created_at = Column(DateTime, default=_utcnow)
+
+
+class CheckIn(Base):
+ __tablename__ = "checkins"
+
+ id = Column(Integer, primary_key=True, autoincrement=True)
+ iso_date = Column(String, nullable=False) # YYYY-MM-DD
+ score = Column(Integer, nullable=False)
+ level = Column(String, nullable=False) # deep/warm/growing/different/learning
+ my_mood = Column(String, nullable=False)
+ partner_guess = Column(String, nullable=False)
+ created_at = Column(DateTime, default=_utcnow)
+
+
+class GardenState(Base):
+ __tablename__ = "garden_state"
+
+ id = Column(Integer, primary_key=True, default=1) # singleton (always id=1)
+ interaction_count = Column(Integer, default=0)
+ total_level = Column(Integer, default=1)
+ last_mood = Column(String, default="happy")
+ updated_at = Column(DateTime, default=_utcnow, onupdate=_utcnow)
diff --git a/src/stage4/requirements.txt b/src/stage4/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..e472edfbec2181aed91ab5e6970b7e9e320ab4d1
--- /dev/null
+++ b/src/stage4/requirements.txt
@@ -0,0 +1,3 @@
+fastapi>=0.110.0
+uvicorn>=0.29.0
+python-multipart>=0.0.9