Update app.py

app.py

@@ -26,7 +26,7 @@ import japanize_matplotlib
 import torch
 import torch.nn as nn
 from huggingface_hub import list_repo_files, hf_hub_download
-from s3prl.nn import S3PRLUpstream
+from s3prl.nn import S3PRLUpstream, Featurizer
 
 # ===== フォント設定 =====
 rcParams["font.family"] = "DejaVu Sans"
@@ -80,67 +80,50 @@ KUSHINADA_REPO = "imprt/kushinada-hubert-base-jtes-er"
 @st.cache_resource(show_spinner=False)
 def load_kushinada_s3prl():
     """
-    S3PRL上流(HuBERT base) + HFの下流(ckpt)を自動取得して復元。
-    - .pt 以外に .ckpt / .pth / .bin も探索
-    - サブフォルダ内も対象
-    - 必要なら KUSHINADA_FILENAME / KUSHINADA_REVISION を Secrets に設定して固定
+    S3PRL 上流(HuBERT base) → Featurizer で [B,T,H] を得る。
+    下流（.ckpt）は線形層の W,b を抽出して SimpleLinearHead を構築。
     """
     token = os.getenv("HF_TOKEN")
     if not token:
         raise RuntimeError("環境変数 HF_TOKEN が見つかりません。SpacesのSettings→Secretsで設定してください。")
 
     revision = os.getenv("KUSHINADA_REVISION", "main")
-    prefer_filename = os.getenv("KUSHINADA_FILENAME")  # 例: "checkpoints/epoch=9-step=1234.ckpt"
+    prefer_filename = os.getenv("KUSHINADA_FILENAME")  # 例: "s3prl/result/downstream/.../dev-best.ckpt"
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    # 1) S3PRL 上流：HuBERT base
+    # 1) 上流 + Featurizer（最終層）
     upstream = S3PRLUpstream("hubert_base").to(device).eval()
+    featurizer = Featurizer(upstream, layer=-1).to(device).eval()
 
-    # 2) モデル内のファイル一覧を取得（サブフォルダ込み）
+    # 2) モデル内のファイル一覧（サブフォルダ込み）
     api = HfApi()
     info = api.model_info(KUSHINADA_REPO, token=token, revision=revision)
-    all_files = [s.rfilename for s in info.siblings]  # ルート/サブフォルダ含むファイルパス
+    all_files = [s.rfilename for s in info.siblings]
 
-    # --- デバッグ出力（必要なら見える化）
-    with st.expander("📦 モデル内ファイル一覧（デバッグ）", expanded=False):
-        st.write(all_files)
-
-    # 3) 候補ファイルの決定
+    # 3) チェックポイント選定
     exts = (".pt", ".ckpt", ".pth", ".bin")
     candidates = [f for f in all_files if f.lower().endswith(exts)]
-
-    # Secretsで明示指定があればそれを優先
     filename = None
     if prefer_filename:
         if prefer_filename in all_files:
             filename = prefer_filename
         else:
-            # サブフォルダなしで指定された場合に補正を試みる
             matches = [f for f in all_files if f.endswith(prefer_filename)]
             if matches:
                 filename = matches[0]
-
-    # それでも未決なら候補の先頭を採用
     if filename is None and candidates:
-        # なるべく "downstream", "classifier", "jtes" を含むものを優先
         ranked = sorted(
             candidates,
             key=lambda f: (
-                -int(any(k in f.lower() for k in ["downstream", "classifier", "jtes", "kushinada"])),
+                -int(any(k in f.lower() for k in ["downstream","classifier","jtes","kushinada"])),
                 len(f)
             )
         )
         filename = ranked[0] if ranked else None
-
     if filename is None:
-        raise FileNotFoundError(
-            "下流チェックポイント(.pt/.ckpt/.pth/.bin)が見つかりません。\n"
-            "モデルページの Files でファイル名を確認し、SpacesのSecretsに "
-            "KUSHINADA_FILENAME として保存してください。"
-        )
+        raise FileNotFoundError("下流チェックポイント(.pt/.ckpt/.pth/.bin)が見つかりません。KUSHINADA_FILENAME を Secrets に指定してください。")
 
-    # 4) チェックポイントを取得
     ckpt_path = hf_hub_download(
         repo_id=KUSHINADA_REPO,
         filename=filename,
@@ -151,18 +134,16 @@ def load_kushinada_s3prl():
         local_dir_use_symlinks=False,
         force_download=False
     )
-
     ckpt = torch.load(ckpt_path, map_location="cpu")
 
-    # 5) state_dict を探索し、線形ヘッド (W, b) を復元
+    # 4) state_dict から線形層の W, b を抽出
     state = None
     if isinstance(ckpt, dict):
-        for key in ["state_dict", "Downstream", "model", "downstream", "net", "weights"]:
+        for key in ["state_dict","Downstream","model","downstream","net","weights"]:
             if key in ckpt and isinstance(ckpt[key], dict):
                 state = ckpt[key]; break
         if state is None:
-            state = ckpt  # そのままstate_dictの場合
-
+            state = ckpt
     if not isinstance(state, dict):
         raise RuntimeError("チェックポイント形式を解釈できませんでした。")
 
@@ -175,26 +156,23 @@ def load_kushinada_s3prl():
                 linear_W, linear_b = v, state[bias_key]
                 break
     if linear_W is None:
-        # weight/biasのペア探索（末尾名が weight/bias）
         twos = [(k,v) for k,v in state.items() if isinstance(v, torch.Tensor) and v.ndim==2 and k.endswith("weight")]
         for wk, w in twos:
-            bk = wk.replace("weight", "bias")
+            bk = wk.replace("weight","bias")
             if bk in state and isinstance(state[bk], torch.Tensor) and state[bk].ndim == 1:
                 linear_W, linear_b = w, state[bk]
                 break
     if linear_W is None:
-        raise RuntimeError("線形分類器の重みが見つかりません。S3PRLのDownstream構造を再現する必要があります。")
+        raise RuntimeError("線形分類器の重みが見つかりません。Downstream 構造の再現が必要です。")
 
     num_classes, hidden_dim = linear_W.shape  # [C, H]
-    head = SimpleLinearHead(in_dim=hidden_dim, num_classes=num_classes,
-                            W=linear_W, b=linear_b).to(device).eval()
+    head = SimpleLinearHead(in_dim=hidden_dim, num_classes=num_classes, W=linear_W, b=linear_b).to(device).eval()
 
-    # ラベル（暫定）。必要なら順序を手動調整してください。
-    default_labels = ["angry", "happy", "neutral", "sad"]
+    default_labels = ["angry","happy","neutral","sad"]
     id2label = {i: (default_labels[i] if num_classes == 4 and i < 4 else f"class_{i}") for i in range(num_classes)}
 
     st.info(f"✅ 使うチェックポイント: `{filename}`（revision: {revision}）")
-    return upstream, head, id2label, device
+    return featurizer, head, id2label, device
 
 # ===== ユーティリティ =====
 def to_wav_bytes(any_bytes: bytes, target_sr=16000, mono=True) -> bytes:
@@ -307,12 +285,11 @@ def _normalize_label(lbl: str) -> str:
 
 def predict_emotion_ai(audio_bytes):
     """
-    S3PRL上流 + HF下流(.ckpt) で推論（S3PRLはリスト入力/出力前提）。
-    入力は CPU の list[Tensor] / list[int] に統一。
-    出力は最終的に list[Tensor([T_i,H])] に正規化 → 時間平均で [B,H] → 線形ヘッド。
+    S3PRL Featurizer で必ず [B,T,H] を取得 → 各サンプルの有効長 reps_len で時間平均 → [B,H]。
+    その後、線形ヘッド（W,b）で分類。
     """
     try:
-        upstream, head, id2label, device = load_kushinada_s3prl()
+        featurizer, head, id2label, device = load_kushinada_s3prl()
     except Exception as e:
         st.error(f"モデルのロードに失敗しました: {e}")
         st.info("音声特徴量ベースの分析に切り替えます。")
@@ -330,94 +307,25 @@ def predict_emotion_ai(audio_bytes):
             y = y[:max_samples]
             st.warning("音声が30秒を超えたため、最初の30秒のみを分析します。")
 
-        # ★ S3PRLは「CPUの list 形式」を想定している実装が多い
-        wavs = [torch.tensor(y, dtype=torch.float32)]   # list[Tensor([T])]
-        wavs_len = [int(len(y))]                        # list[int]
+        # S3PRLは list[Tensor], list[int] を想定
+        wavs = [torch.tensor(y, dtype=torch.float32)]
+        wavs_len = [int(len(y))]
 
         with torch.no_grad():
-            # 返り値は実装により:
-            #  - (list[Tensor([T_i,H])], list[int])  ←最も一般的
-            #  - Tensor([B,T,H]) / dict / 入れ子
-            reps_out = upstream(wavs, wavs_len)
-
-            def as_seq_list(obj):
-                """上流出力を list[Tensor([T_i, H])] に正規化"""
-                # 1) すでに (seqs, lens) 形式
-                if isinstance(obj, tuple) and len(obj) == 2:
-                    seqs, lens = obj
-                    # seqs: list of Tensor ならそのまま
-                    if isinstance(seqs, list) and len(seqs) > 0 and isinstance(seqs[0], torch.Tensor):
-                        return seqs
-                    # seqs が Tensor [B,T,H] の場合 → バラす
-                    if isinstance(seqs, torch.Tensor):
-                        if seqs.dim() == 3:
-                            return [seqs[i].cpu() for i in range(seqs.size(0))]
-                        if seqs.dim() == 2:
-                            return [seqs.cpu()]
-                    # dict などが来たら再帰
-                    return as_seq_list(seqs)
-
-                # 2) Tensor
-                if isinstance(obj, torch.Tensor):
-                    if obj.dim() == 3:    # [B,T,H]
-                        return [obj[i].cpu() for i in range(obj.size(0))]
-                    if obj.dim() == 2:    # [T,H]
-                        return [obj.cpu()]
-                    if obj.dim() == 1:    # [H]（既にプール済み）→T=1として扱う
-                        return [obj.unsqueeze(0).cpu()]
-
-                # 3) dict（代表キー優先）
-                if isinstance(obj, dict):
-                    for k in ["last_hidden_state", "hidden_states"]:
-                        if k in obj:
-                            v = obj[k]
-                            # hidden_states がリストなら最終層
-                            if k == "hidden_states" and isinstance(v, (list, tuple)) and len(v) > 0:
-                                v = v[-1]
-                            return as_seq_list(v)
-                    # 他キーも探索
-                    for v in obj.values():
-                        got = as_seq_list(v)
-                        if got:
-                            return got
-                    return []
-
-                # 4) list / tuple（入れ子を平坦化）
-                if isinstance(obj, (list, tuple)):
-                    out = []
-                    for it in obj:
-                        out.extend(as_seq_list(it))
-                    return out
-
-                # それ以外は無視
-                return []
-
-            seq_list = as_seq_list(reps_out)
-            if not seq_list:
-                raise RuntimeError("上流出力を [T,H] の列へ正規化できませんでした。")
-
-            # ★ 時間平均で [H] にプール → [B,H]
-            pooled_list = []
-            for t in seq_list:
-                if not isinstance(t, torch.Tensor):
-                    continue
-                t = t.to(device)
-                if t.dim() == 3:      # [?,T,H] が来たら T 次元で平均
-                    t = t.mean(dim=1)
-                if t.dim() == 2:      # [T,H]
-                    pooled_list.append(t.mean(dim=0))   # -> [H]
-                elif t.dim() == 1:    # [H]
-                    pooled_list.append(t)
-                else:
-                    raise RuntimeError(f"Unexpected tensor shape from upstream: {tuple(t.size())}")
-
-            if len(pooled_list) == 0:
-                raise RuntimeError("プーリング後テンソルが空です。")
-
-            pooled = torch.stack(pooled_list, dim=0)    # [B,H]
-
-            # 線形ヘッドで分類
-            logits = head.fc(pooled)                     # [B,C]
+            reps, reps_len = featurizer(wavs, wavs_len)  # reps: [B,T,H], reps_len: list[int] or Tensor[B]
+            if isinstance(reps_len, torch.Tensor):
+                reps_len = reps_len.tolist()
+
+            # 有効長のみで平均（パディングを無視）
+            pooled = []
+            for i in range(reps.shape[0]):
+                Ti = int(reps_len[i]) if reps_len else reps.shape[1]
+                Ti = max(1, min(Ti, reps.shape[1]))  # 安全側
+                pooled.append(reps[i, :Ti].mean(dim=0))
+            pooled = torch.stack(pooled, dim=0)  # [B,H]
+
+            # 線形ヘッドで分類（head.fcに直接入れる）
+            logits = head.fc(pooled.to(device))  # [B,C]
             probs = torch.softmax(logits, dim=-1)[0].detach().cpu().numpy()
 
         pred_id = int(np.argmax(probs))
@@ -435,7 +343,6 @@ def predict_emotion_ai(audio_bytes):
         return label, scores, "AI(S3PRL)"
 
     except Exception as e:
-        # デバッグ補助（発生時だけ型を少し表示）
         st.warning(f"AI予測中にエラーが発生: {e}")
         return predict_emotion_features(audio_bytes)