{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":102335,"databundleVersionId":12518947,"sourceType":"competition"},{"sourceId":13468194,"sourceType":"datasetVersion","datasetId":8520152}],"dockerImageVersionId":31089,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport pathlib\nimport re\nimport math\nimport json\nimport numpy as np\nimport pandas as pd\nimport polars as pl\nimport joblib\nimport warnings\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\n\nfrom scipy.spatial.transform import Rotation as R\nimport kaggle_evaluation.cmi_inference_server\n\nwarnings.filterwarnings(\"ignore\")\n\n# =============================================================================\n# 1) CONFIG\n# =============================================================================\nclass Config:\n    DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    ACC_COLS = [\"acc_x\", \"acc_y\", \"acc_z\"]\n    ROT_COLS = [\"rot_w\", \"rot_x\", \"rot_y\", \"rot_z\"]\n    BATCH_SIZE = 1\n\n# =============================================================================\n# 2) PREPROCESS & FEATURES  \n# =============================================================================\ndef handle_quaternion_missing_values(rot_data: np.ndarray) -> np.ndarray:\n    rot_cleaned = rot_data.copy()\n    for i in range(len(rot_data)):\n        row = rot_data[i]\n        missing_count = np.isnan(row).sum()\n        if missing_count == 0:\n            norm = np.linalg.norm(row)\n            rot_cleaned[i] = row / norm if norm > 1e-8 else [1.0, 0.0, 0.0, 0.0]\n        elif missing_count == 1:\n            missing_idx = np.where(np.isnan(row))[0][0]\n            valid_values = row[~np.isnan(row)]\n            sum_squares = np.sum(valid_values**2)\n            if sum_squares <= 1.0:\n                missing_value = np.sqrt(max(0, 1.0 - sum_squares))\n                if i > 0 and not np.isnan(rot_cleaned[i-1, missing_idx]) and rot_cleaned[i-1, missing_idx] < 0:\n                    missing_value = -missing_value\n                rot_cleaned[i, missing_idx] = missing_value\n            else:\n                rot_cleaned[i] = [1.0, 0.0, 0.0, 0.0]\n        else:\n            rot_cleaned[i] = [1.0, 0.0, 0.0, 0.0]\n    return rot_cleaned\n\ndef calculate_angular_velocity_from_quat(rot_data, time_delta=1/200):\n    quat_values = rot_data[:, [1, 2, 3, 0]] if rot_data.shape[1] == 4 else rot_data\n    num_samples = quat_values.shape[0]\n    angular_vel = np.zeros((num_samples, 3))\n    for i in range(num_samples - 1):\n        q_t, q_t_plus_dt = quat_values[i], quat_values[i+1]\n        if np.any(np.isnan(q_t)) or np.any(np.isnan(q_t_plus_dt)): \n            continue\n        try:\n            rot_t, rot_t_plus_dt = R.from_quat(q_t), R.from_quat(q_t_plus_dt)\n            delta_rot = rot_t.inv() * rot_t_plus_dt\n            angular_vel[i, :] = delta_rot.as_rotvec() / time_delta\n        except ValueError:\n            pass\n    return angular_vel\n\ndef calculate_angular_distance(rot_data):\n    quat_values = rot_data[:, [1, 2, 3, 0]] if rot_data.shape[1] == 4 else rot_data\n    num_samples = quat_values.shape[0]\n    angular_dist = np.zeros(num_samples)\n    for i in range(num_samples - 1):\n        q1, q2 = quat_values[i], quat_values[i+1]\n        if np.any(np.isnan(q1)) or np.any(np.isnan(q2)): \n            continue\n        try:\n            r1, r2 = R.from_quat(q1), R.from_quat(q2)\n            angle = np.linalg.norm((r1.inv() * r2).as_rotvec())\n            angular_dist[i] = angle\n        except ValueError:\n            pass\n    return angular_dist\n\ndef remove_gravity_from_acc(acc_data, rot_data):\n    acc_values = acc_data.values if isinstance(acc_data, pd.DataFrame) else acc_data\n    quat_values = rot_data.values if isinstance(rot_data, pd.DataFrame) else rot_data\n    quat_scipy = quat_values[:, [1, 2, 3, 0]]\n    num_samples = acc_values.shape[0]\n    linear_accel = np.zeros_like(acc_values)\n    gravity_world = np.array([0, 0, 9.81])\n    for i in range(num_samples):\n        if np.any(np.isnan(quat_scipy[i])):\n            linear_accel[i, :] = acc_values[i, :]\n            continue\n        try:\n            rotation = R.from_quat(quat_scipy[i])\n            gravity_sensor_frame = rotation.apply(gravity_world, inverse=True)\n            linear_accel[i, :] = acc_values[i, :] - gravity_sensor_frame\n        except ValueError:\n            linear_accel[i, :] = acc_values[i, :]\n    return linear_accel\n\ndef add_features(df: pd.DataFrame) -> pd.DataFrame:\n    # TÍNH THEO NHÓM sequence_id giống bản train\n    df[\"acc_mag\"] = np.linalg.norm(df[Config.ACC_COLS].values, axis=1)\n    df[\"acc_mag_jerk\"] = df.groupby(\"sequence_id\")[\"acc_mag\"].diff().fillna(0)\n    df[\"jerk_x\"], df[\"jerk_y\"], df[\"jerk_z\"] = np.gradient(df[\"acc_x\"]), np.gradient(df[\"acc_y\"]), np.gradient(df[\"acc_z\"])\n    df[\"jerk_magnitude\"] = np.linalg.norm(df[[\"jerk_x\", \"jerk_y\", \"jerk_z\"]].values, axis=1)\n\n    window = 20\n    for _, g in df.groupby(\"sequence_id\"):\n        df.loc[g.index, \"acc_xy_corr\"] = g[\"acc_x\"].rolling(window, min_periods=1).corr(g[\"acc_y\"]).fillna(0)\n        df.loc[g.index, \"acc_xz_corr\"] = g[\"acc_x\"].rolling(window, min_periods=1).corr(g[\"acc_z\"]).fillna(0)\n        df.loc[g.index, \"acc_yz_corr\"] = g[\"acc_y\"].rolling(window, min_periods=1).corr(g[\"acc_z\"]).fillna(0)\n\n    df[\"rot_angle\"] = 2 * np.arccos(df[\"rot_w\"].clip(-1, 1))\n    df[\"rot_angle_vel\"] = df.groupby(\"sequence_id\")[\"rot_angle\"].diff().fillna(0)\n\n    rot_numpy = df[Config.ROT_COLS].to_numpy()\n    angular_vel = calculate_angular_velocity_from_quat(rot_numpy)\n    angular_dist = calculate_angular_distance(rot_numpy)\n    df[[\"angular_vel_x\", \"angular_vel_y\", \"angular_vel_z\"]] = angular_vel\n    df[\"angular_distance\"] = angular_dist\n    df[\"angular_vel_magnitude\"] = np.linalg.norm(angular_vel, axis=1)\n\n    linear_accel = remove_gravity_from_acc(df[Config.ACC_COLS], df[Config.ROT_COLS])\n    df[[\"acc_x2\", \"acc_y2\", \"acc_z2\"]] = linear_accel\n    df[\"acc_mag2\"] = np.linalg.norm(linear_accel, axis=1)\n    df[\"acc_mag_jerk2\"] = df.groupby(\"sequence_id\")[\"acc_mag2\"].diff().fillna(0)\n    df[\"jerk_x2\"], df[\"jerk_y2\"], df[\"jerk_z2\"] = np.gradient(df[\"acc_x2\"]), np.gradient(df[\"acc_y2\"]), np.gradient(df[\"acc_z2\"])\n    df[\"jerk_magnitude2\"] = np.linalg.norm(df[[\"jerk_x2\", \"jerk_y2\", \"jerk_z2\"]].values, axis=1)\n\n    for _, g in df.groupby(\"sequence_id\"):\n        df.loc[g.index, \"acc_xy_corr2\"] = g[\"acc_x2\"].rolling(window, min_periods=1).corr(g[\"acc_y2\"]).fillna(0)\n        df.loc[g.index, \"acc_xz_corr2\"] = g[\"acc_x2\"].rolling(window, min_periods=1).corr(g[\"acc_z2\"]).fillna(0)\n        df.loc[g.index, \"acc_yz_corr2\"] = g[\"acc_y2\"].rolling(window, min_periods=1).corr(g[\"acc_z2\"]).fillna(0)\n\n    df.replace([np.inf, -np.inf], 0, inplace=True)\n    df.fillna(0, inplace=True)\n    return df\n\ndef preprocess_left_handed_pd(df: pd.DataFrame, demo: pd.DataFrame) -> pd.DataFrame:\n    pl_df = pl.DataFrame(df)\n    pl_demo = pl.DataFrame(demo)\n    l_subjs = pl_demo.filter(pl.col(\"handedness\") == 0)[\"subject\"].to_list()\n\n    r_tr = pl_df.filter(~pl.col(\"subject\").is_in(l_subjs))\n    l_tr = pl_df.filter(pl.col(\"subject\").is_in(l_subjs))\n    if l_tr.shape[0] == 0:\n        return r_tr.to_pandas()\n\n    # Flip trục y, z gia tốc; mirror quaternion như bản train\n    l_tr = l_tr.with_columns(-pl.col(\"acc_y\"))\n    l_tr = l_tr.with_columns(-pl.col(\"acc_z\"))\n\n    rot_np = l_tr.select(Config.ROT_COLS).to_numpy()\n    rot_scipy = rot_np[:, [1, 2, 3, 0]]\n    r = R.from_quat(rot_scipy)\n    euler = r.as_euler(\"xyz\")\n    euler[:, 1] = -euler[:, 1]\n    euler[:, 2] = -euler[:, 2]\n    r = R.from_euler(\"xyz\", euler)\n    q = r.as_quat()\n    l_tr = l_tr.with_columns(pl.DataFrame(q, schema=[\"rot_x\", \"rot_y\", \"rot_z\", \"rot_w\"]))\n\n    pl_df2 = pl.concat([r_tr, l_tr]).sort(by=\"row_id\")\n    return pl_df2.to_pandas()\n\ndef pad_sequences(sequences: list, maxlen: int, padding: str=\"pre\", truncating: str=\"pre\", dtype: str=\"float32\") -> np.ndarray:\n    n_samples = len(sequences); sample_shape = tuple()\n    for s in sequences:\n        if len(s) > 0:\n            sample_shape = np.asarray(s).shape[1:]; break\n    x = np.zeros((n_samples, maxlen) + sample_shape, dtype=dtype)\n    for idx, s in enumerate(sequences):\n        if len(s) == 0: continue\n        s_np = np.asarray(s, dtype=dtype)\n        if truncating == \"pre\": s_np = s_np[-maxlen:]\n        else: s_np = s_np[:maxlen]\n        trunc = s_np.shape[0]\n        if padding == \"pre\": x[idx, -trunc:] = s_np\n        else: x[idx, :trunc] = s_np\n    return x\n\n# --- Pyramid Temporal Feature (PTF) ---\nclass PyramidTemporalFeature(nn.Module):\n    \"\"\"\n    Module trích xuất đặc trưng theo nhiều tần suất thời gian khác nhau.\n    Tương tự FPN nhưng dùng nhiều Conv1D kernel size khác nhau để nắm bắt pattern ngắn & dài.\n    \"\"\"\n    def __init__(self, in_channels, out_channels, kernel_sizes=[3, 5, 7]):\n        super().__init__()\n        self.branches = nn.ModuleList([\n            nn.Sequential(\n                nn.Conv1d(in_channels, out_channels, kernel_size=k, padding=k//2),\n                nn.BatchNorm1d(out_channels),\n                nn.ReLU(inplace=True)\n            )\n            for k in kernel_sizes\n        ])\n        # Gộp lại các scale\n        self.fuse = nn.Conv1d(out_channels * len(kernel_sizes), out_channels, kernel_size=1)\n\n    def forward(self, x):\n        # x: [B, C, T]\n        features = [branch(x) for branch in self.branches]  # nhiều scale thời gian\n        concat = torch.cat(features, dim=1)  # ghép theo kênh\n        out = self.fuse(concat)              # hợp nhất lại\n        return out                           # [B, out_channels, T]\n\n\n# =============================================================================\n# 3) MODEL (giống bản train)\n# =============================================================================\nclass SEBlock(nn.Module):\n    def __init__(self, channel, reduction=8):\n        super().__init__()\n        self.avg_pool = nn.AdaptiveAvgPool1d(1)\n        self.fc = nn.Sequential(\n            nn.Linear(channel, channel // reduction), nn.ReLU(inplace=True),\n            nn.Linear(channel // reduction, channel), nn.Sigmoid())\n    def forward(self, x):\n        b, c, _ = x.size()\n        y = self.avg_pool(x).view(b, c)\n        y = self.fc(y).view(b, c, 1)\n        return x * y.expand_as(x)\n\nclass ResidualSEBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, kernel_size, drop=0.3):\n        super().__init__()\n        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size, padding=\"same\", bias=False)\n        self.bn1 = nn.BatchNorm1d(out_channels)\n        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size, padding=\"same\", bias=False)\n        self.bn2 = nn.BatchNorm1d(out_channels)\n        self.se = SEBlock(out_channels)\n        self.relu = nn.ReLU()\n        self.dropout = nn.Dropout(drop)\n        self.pool = nn.MaxPool1d(2)\n        self.shortcut = nn.Sequential(\n            nn.Conv1d(in_channels, out_channels, 1, bias=False),\n            nn.BatchNorm1d(out_channels)\n        ) if in_channels != out_channels else nn.Identity()\n    def forward(self, x):\n        identity = self.shortcut(x)\n        out = self.relu(self.bn1(self.conv1(x)))\n        out = self.se(self.bn2(self.conv2(out)))\n        out = self.dropout(self.pool(self.relu(out + identity)))\n        return out\n\nclass AttentionLayer(nn.Module):\n    def __init__(self, hidden_dim):\n        super().__init__()\n        self.attn = nn.Sequential(nn.Linear(hidden_dim, 1), nn.Tanh())\n\n    def forward(self, x):\n        weights = F.softmax(self.attn(x).squeeze(-1), dim=1).unsqueeze(-1)\n        return torch.sum(x * weights, dim=1)\n\n\nclass CrossAttention(nn.Module):\n    def __init__(self, feature_dim, num_heads=8, dropout=0.1):\n        super().__init__()\n        self.feature_dim = feature_dim\n        self.num_heads = num_heads\n        self.head_dim = feature_dim // num_heads\n        assert feature_dim % num_heads == 0, \"feature_dim must be divisible by num_heads\"\n        self.q_linear = nn.Linear(feature_dim, feature_dim, bias=False)\n        self.k_linear = nn.Linear(feature_dim, feature_dim, bias=False)\n        self.v_linear = nn.Linear(feature_dim, feature_dim, bias=False)\n        self.out_linear = nn.Linear(feature_dim, feature_dim)\n        self.dropout = nn.Dropout(dropout)\n        self.layer_norm = nn.LayerNorm(feature_dim)\n    def forward(self, query_branch, key_value_branches):\n        B, T, C = query_branch.shape\n        Q = self.q_linear(query_branch)\n        all_kv = torch.cat(key_value_branches, dim=1)\n        K = self.k_linear(all_kv)\n        V = self.v_linear(all_kv)\n        Q = Q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)\n        K = K.view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)\n        V = V.view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)\n        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)\n        attn_weights = F.softmax(scores, dim=-1)\n        attn_weights = self.dropout(attn_weights)\n        attn_output = torch.matmul(attn_weights, V)\n        attn_output = attn_output.transpose(1, 2).contiguous().view(B, T, C)\n        output = self.out_linear(attn_output)\n        output = self.layer_norm(query_branch + output)\n        return output\n\nclass IMUCrossAttentionFusion(nn.Module):\n    def __init__(self, feature_dim=256, num_heads=8, dropout=0.1):\n        super().__init__()\n        self.cross_attn1 = CrossAttention(feature_dim, num_heads, dropout)\n        self.cross_attn2 = CrossAttention(feature_dim, num_heads, dropout)\n        self.cross_attn3 = CrossAttention(feature_dim, num_heads, dropout)\n    def forward(self, imu1, imu2, imu3):\n        e1 = self.cross_attn1(imu1, [imu2, imu3])\n        e2 = self.cross_attn2(imu2, [imu1, imu3])\n        e3 = self.cross_attn3(imu3, [imu1, imu2])\n        return e1, e2, e3\n\nclass IMUCrossAttentionModel_PTF(nn.Module):\n    def __init__(self, imu_dim, n_classes):\n        super().__init__()\n        self.imu_dim = imu_dim\n\n        # --- 3 nhánh CNN backbone cho dữ liệu IMU ---\n        self.imu_branch1 = nn.Sequential(\n            ResidualSEBlock(12, 128, 3),\n            ResidualSEBlock(128, 256, 5)\n        )\n        self.imu_branch2 = nn.Sequential(\n            ResidualSEBlock(11, 128, 3),\n            ResidualSEBlock(128, 256, 5)\n        )\n        self.imu_branch3 = nn.Sequential(\n            ResidualSEBlock(12, 128, 3),\n            ResidualSEBlock(128, 256, 5)\n        )\n\n        # --- Pyramid Temporal Feature cho mỗi nhánh ---\n        self.ptf1 = PyramidTemporalFeature(256, 128)\n        self.ptf2 = PyramidTemporalFeature(256, 128)\n        self.ptf3 = PyramidTemporalFeature(256, 128)\n\n        # --- Cross attention fusion ---\n        self.cross_attention_fusion = IMUCrossAttentionFusion(feature_dim=128)\n\n        # --- BiLSTM + attention + FC ---\n        self.bilstm = nn.LSTM(128 * 3, 512, bidirectional=True, batch_first=True)\n        self.attention = AttentionLayer(1024)\n        self.fc = nn.Linear(1024, n_classes)\n    \n    def forward(self, x):\n        imu = x[:, :, :self.imu_dim]\n\n        # tách dữ liệu IMU cho 3 nhánh (tuỳ theo cấu trúc cảm biến)\n        imu1 = self.imu_branch1(imu[:, :, :12].transpose(1, 2))     # [B, C, T]\n        imu2 = self.imu_branch2(imu[:, :, 12:23].transpose(1, 2))\n        imu3 = self.imu_branch3(imu[:, :, 23:].transpose(1, 2))\n        \n        # --- Trích đặc trưng đa tần số thời gian ---\n        imu1_ptf = self.ptf1(imu1).transpose(1, 2)\n        imu2_ptf = self.ptf2(imu2).transpose(1, 2)\n        imu3_ptf = self.ptf3(imu3).transpose(1, 2)\n\n        # --- Cross attention fusion giữa 3 nhánh ---\n        imu1, imu2, imu3 = self.cross_attention_fusion(imu1_ptf, imu2_ptf, imu3_ptf)\n    \n        # --- Ghép & đưa qua BiLSTM ---\n        merged = torch.cat((imu1, imu2, imu3), dim=2)\n        lstm_out, _ = self.bilstm(merged)\n\n        # --- Attention + phân lớp ---\n        attended = self.attention(lstm_out)\n        return self.fc(attended)\n\n# =============================================================================\n# 4) ENS INFERENCE UTILITIES\n# =============================================================================\nclass CMIInferenceDataset(Dataset):\n    def __init__(self, data_array):\n        self.data_array = data_array\n    def __len__(self):\n        return 1\n    def __getitem__(self, idx):\n        return torch.FloatTensor(self.data_array)\n\ndef _find_model_dir_and_folds():\n    \"\"\"\n    Tìm thư mục chứa artifacts. Ưu tiên /kaggle/working, sau đó các input dir.\n    Kỳ vọng tồn tại các cặp:\n      - prep_fold_{k}.joblib\n      - model_best_fold_{k}.pt\n    Fallback: nếu không có theo fold, dùng single model 'prep.joblib' + 'model_best.pt'.\n    \"\"\"\n    candidates = [pathlib.Path(\"/kaggle/working\"), pathlib.Path(\"/kaggle/input/lstm-triet\")]\n    for base in candidates:\n        if not base.exists():\n            continue\n        fold_files = sorted(base.glob(\"prep_fold_*.joblib\"))\n        model_files = sorted(base.glob(\"model_best_fold_*.pt\"))\n        if fold_files and model_files:\n            # lấy list fold chung giữa 2 bên\n            folds = []\n            for p in fold_files:\n                m = re.search(r\"prep_fold_(\\d+)\\.joblib$\", p.name)\n                if m:\n                    k = int(m.group(1))\n                    if (base / f\"model_best_fold_{k}.pt\").exists():\n                        folds.append(k)\n            folds = sorted(set(folds))\n            if folds:\n                return base, folds, True\n        # fallback single\n        if (base / \"prep.joblib\").exists() and (base / \"model_best.pt\").exists():\n            return base, [], False\n    raise FileNotFoundError(\"Không tìm thấy artifacts theo fold hoặc single trong các thư mục dự kiến.\")\n\ndef _prepare_sequence(sequence_df: pd.DataFrame, demo_df: pd.DataFrame) -> pd.DataFrame:\n    # xử lý giống train\n    sequence_df = sequence_df.copy()\n    sequence_df[Config.ROT_COLS] = handle_quaternion_missing_values(sequence_df[Config.ROT_COLS].to_numpy())\n    sequence_df = pd.concat([g.ffill().bfill().fillna(0) for _, g in sequence_df.groupby(\"sequence_id\")], axis=0)\n    sequence_df = preprocess_left_handed_pd(sequence_df, demo_df)\n    sequence_df = add_features(sequence_df)\n    return sequence_df\n\ndef _build_features(sequence_df: pd.DataFrame, features: list, transformer, max_length: int) -> np.ndarray:\n    seq_X = sequence_df[features].to_numpy()\n    # seq_X = transformer.transform(seq_X)\n    if len(seq_X) > max_length:\n        padded = seq_X[-max_length:]\n    else:\n        padded = np.zeros((max_length, seq_X.shape[1]), dtype=np.float32)\n        padded[-len(seq_X):, :] = seq_X\n    return padded\n\ndef create_prediction_function_ensemble(models, preps, categories):\n    \"\"\"\n    models: list[nn.Module]\n    preps : list[dict] (transformer, features, max_length, fold)\n    categories: list[str]\n    \"\"\"\n    def predict(sequence: pl.DataFrame, demographics: pl.DataFrame) -> str:\n        try:\n            seq_df = sequence.to_pandas()\n            demo_df = demographics.to_pandas()\n\n            required_cols = set(Config.ACC_COLS + Config.ROT_COLS + [\"sequence_id\", \"row_id\", \"subject\"])\n            missing = required_cols - set(seq_df.columns)\n            if missing:\n                raise ValueError(f\"Thiếu cột bắt buộc: {missing}\")\n\n            seq_df = _prepare_sequence(seq_df, demo_df)\n\n            probs_all = []\n            for model, prep in zip(models, preps):\n                feats = prep[\"features\"]\n                X_pad = _build_features(seq_df, feats, prep[\"transformer\"], prep[\"max_length\"])\n                ds = CMIInferenceDataset(X_pad)\n                loader = DataLoader(ds, batch_size=Config.BATCH_SIZE)\n                model.eval()\n                with torch.no_grad():\n                    for batch_X in loader:\n                        logits = model(batch_X.to(Config.DEVICE))\n                        probs = torch.softmax(logits, dim=1).cpu().numpy()\n                        probs_all.append(probs[0])  # (C,)\n            probs_mean = np.mean(probs_all, axis=0)\n            pred_idx = int(np.argmax(probs_mean))\n            return categories[pred_idx]\n        except Exception as e:\n            print(f\"[Predict Error] {e}\")\n            return categories[0]\n    return predict\n\n# =============================================================================\n# 5) MAIN: LOAD ENSEMBLE hoặc FALLBACK SINGLE\n# =============================================================================\ndef main():\n    model_dir, folds, is_fold = _find_model_dir_and_folds()\n    print(f\"Artifacts dir: {model_dir} | by_fold={is_fold} | folds={folds}\")\n\n    if is_fold:\n        # kiểm tra categories giữa các fold nhất quán\n        preps = []\n        models = []\n        cats_ref = None\n        for k in folds:\n            prep = joblib.load(model_dir / f\"prep_fold_{k}.joblib\")\n            preps.append(prep)\n            if cats_ref is None:\n                cats_ref = prep[\"categories\"]\n            else:\n                assert cats_ref == prep[\"categories\"], \"Mismatch categories giữa các fold!\"\n            model = IMUCrossAttentionModel_PTF(\n                imu_dim=len(prep[\"features\"]),\n                n_classes=len(prep[\"categories\"])\n            ).to(Config.DEVICE)\n            state = torch.load(model_dir / f\"model_best_fold_{k}.pt\", map_location=Config.DEVICE)\n            model.load_state_dict(state)\n            models.append(model)\n        print(f\"✓ Loaded {len(models)} fold models for ensemble.\")\n        predict_fn = create_prediction_function_ensemble(models, preps, cats_ref)\n        return predict_fn\n    else:\n     \n        artifacts = joblib.load(model_dir / \"prep.joblib\")\n        features = artifacts['features']\n        n_classes = len(artifacts['categories'])\n        model = IMUCrossAttentionModel_PTF(imu_dim=len(features), n_classes=n_classes).to(Config.DEVICE)\n        model.load_state_dict(torch.load(model_dir / \"model_best.pt\", map_location=Config.DEVICE))\n        print(\"✓ Loaded single model artifacts.\")\n\n        def predict(sequence: pl.DataFrame, demographics: pl.DataFrame) -> str:\n            try:\n                seq_df = _prepare_sequence(sequence.to_pandas(), demographics.to_pandas())\n                X_pad = _build_features(seq_df, features, artifacts[\"transformer\"], artifacts[\"max_length\"])\n                ds = CMIInferenceDataset(X_pad)\n                loader = DataLoader(ds, batch_size=Config.BATCH_SIZE)\n                model.eval()\n                with torch.no_grad():\n                    for batch_X in loader:\n                        logits = model(batch_X.to(Config.DEVICE))\n                        probs = torch.softmax(logits, dim=1).cpu().numpy()[0]\n                        pred_idx = int(np.argmax(probs))\n                return artifacts[\"categories\"][pred_idx]\n            except Exception as e:\n                print(f\"[Predict Error] {e}\")\n                return artifacts[\"categories\"][0]\n        return predict\n\nif __name__ == '__main__':\n    predict_function = main()\n    inference_server = kaggle_evaluation.cmi_inference_server.CMIInferenceServer(predict_function)\n\n    if os.getenv('KAGGLE_IS_COMPETITION_RERUN'):\n        inference_server.serve()\n    else:\n        test_csv_path = pathlib.Path(\"/kaggle/input/cmi-detect-behavior-with-sensor-data/test.csv\")\n        test_demo_path = pathlib.Path(\"/kaggle/input/cmi-detect-behavior-with-sensor-data/test_demographics.csv\")\n        if test_csv_path.exists():\n            inference_server.run_local_gateway(data_paths=(test_csv_path, test_demo_path))\n        else:\n            print(\"Local test data not found.\")\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-10-22T14:52:15.031920Z","iopub.execute_input":"2025-10-22T14:52:15.032338Z","iopub.status.idle":"2025-10-22T14:52:17.639146Z","shell.execute_reply.started":"2025-10-22T14:52:15.032246Z","shell.execute_reply":"2025-10-22T14:52:17.637195Z"}},"outputs":[],"execution_count":null}]}