{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"},{"sourceId":14165809,"sourceType":"datasetVersion","datasetId":8701543,"isSourceIdPinned":true}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### CONFIG","metadata":{}},{"cell_type":"code","source":"!pip install torch_geometric","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T13:06:35.025509Z","iopub.execute_input":"2025-12-18T13:06:35.025826Z","iopub.status.idle":"2025-12-18T13:06:42.369238Z","shell.execute_reply.started":"2025-12-18T13:06:35.025792Z","shell.execute_reply":"2025-12-18T13:06:42.368284Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport sys\nimport json\nimport numpy as np\nimport torch\nimport pandas as pd\nfrom types import SimpleNamespace\n\ncfg = SimpleNamespace(**{})\ncfg.debug = True\n\n#paths\ncfg.output_dir = f\"/kaggle/working/\"\ncfg.data_folder = f\"/kaggle/input/MABe-mouse-behavior-detection\"\ncfg.train_df = f'/kaggle/input/MABe-mouse-behavior-detection/test.csv'\n\n# stages\ncfg.test = False\ncfg.train = True\ncfg.train_val =  False\ncfg.eval_epochs = 1\ncfg.seed = 1994\n\n    \n# Sampling strategy for long videos\ncfg.min_windows_per_video = 5  # Minimum windows to sample per video\ncfg.max_windows_per_video = 1000 # Maximum windows to sample per video\ncfg.windows_per_epoch_ratio = 1.0 # Sample 30% of possible windows per epoch\n\ncfg.MASTER_SKELETON = [\n    'nose',\n    'ear_left',\n    'ear_right',\n    'head_center',\n    'body_center',\n    'tail_base'\n]\ncfg.NUM_MASTER_KEYPOINTS = len(cfg.MASTER_SKELETON)\ncfg.MASTER_SKELETON_MAP = {name: i for i, name in enumerate(cfg.MASTER_SKELETON)}\n\ncfg.tracked_bodyparts = ['body_center', 'ear_left', 'ear_right', 'forepaw_left', 'forepaw_right', \n    'head', 'headpiece_bottombackleft', 'headpiece_bottombackright', \n    'headpiece_bottomfrontleft', 'headpiece_bottomfrontright', \n    'headpiece_topbackleft', 'headpiece_topbackright', 'headpiece_topfrontleft',\n    'headpiece_topfrontright', 'hindpaw_left', 'hindpaw_right', 'hip_left',\n    'hip_right', 'lateral_left', 'lateral_right', 'neck', 'nose', 'spine_1', \n    'spine_2', 'tail_base', 'tail_middle_1', 'tail_middle_2', 'tail_midpoint', 'tail_tip'\n]\n\ncfg.set_behavior_classes = ['allogroom', 'approach', 'attack', 'attemptmount', 'avoid', 'biteobject',\n    'chase', 'chaseattack', 'climb', 'defend', 'dig', 'disengage', 'dominance', 'dominancegroom', \n    'dominancemount', 'ejaculate', 'escape', 'exploreobject', 'flinch', 'follow', 'freeze', 'genitalgroom',\n    'huddle', 'intromit', 'mount', 'rear', 'reciprocalsniff', 'rest', 'run', 'self', 'selfgroom', \n    'shepherd', 'sniff', 'sniffbody', 'sniffface', 'sniffgenital', 'submit', 'tussle', 'no_action'\n]\n\ncfg.set_mice = ['mouse1', 'mouse2', 'mouse3', 'mouse4']\ncfg.max_pairs = len(cfg.set_mice) * len(cfg.set_mice)  # 16 if 4 mice\n\ncfg.mouse_id_map = {1:0, 2:1, 3:2, 4:3}\ncfg.mouse_id_to_string = {0:'mouse1', 1:'mouse2', 2:'mouse3', 3:'mouse4'}\n\ncfg.action_id_map = {'allogroom': 0, 'approach': 1, 'attack': 2, 'attemptmount': 3, 'avoid': 4,\n                     'biteobject': 5, 'chase': 6, 'chaseattack': 7, 'climb': 8, 'defend': 9, 'dig': 10,\n                     'disengage': 11, 'dominance': 12, 'dominancegroom': 13, 'dominancemount': 14,\n                     'ejaculate': 15, 'escape': 16, 'exploreobject': 17, 'flinch': 18, 'follow': 19,\n                     'freeze': 20, 'genitalgroom': 21, 'huddle': 22, 'intromit': 23, 'mount': 24,\n                     'rear': 25, 'reciprocalsniff': 26, 'rest': 27, 'run': 28, 'self': 29,\n                     'selfgroom': 30, 'shepherd': 31, 'sniff': 32, 'sniffbody': 33, 'sniffface': 34,\n                     'sniffgenital': 35, 'submit': 36, 'tussle': 37, 'no_action': 38}\n\ncfg.id_to_action_map = {v:k for k,v in cfg.action_id_map.items()}\n\n# DETECTED FROM ACTUAL DATASET - actions that only occur when agent_id == target_id\ncfg.self_only_actions = ['biteobject', 'climb', 'dig', 'exploreobject', 'freeze', 'genitalgroom', 'huddle', 'rear', 'rest', 'run', 'selfgroom']\n\n# DETECTED FROM ACTUAL DATASET - actions that only occur when agent_id != target_id\ncfg.social_only_actions = ['allogroom', 'approach', 'attack', 'attemptmount', 'avoid', 'chase', 'chaseattack', \n                           'defend', 'disengage', 'dominance', 'dominancegroom', 'dominancemount', 'ejaculate', \n                           'escape', 'flinch', 'follow', 'intromit', 'mount', 'reciprocalsniff', 'shepherd',\n                           'sniff', 'sniffbody', 'sniffface', 'sniffgenital', 'submit', 'tussle']\ncfg.reverse_time = False\n\ncfg.window_size = 512\ncfg.feature_dim = 232 #280\ncfg.per_mouse_feature_dim = cfg.feature_dim // len(cfg.set_mice)\ncfg.bias_prob = 0.5  # Probability of sampling behavior-rich window during training\ncfg.oversample_factor = 10\ncfg.preprocessing_basedir = \"/kaggle/working/preprocessed\"\n\n#model\ncfg.model = \"mdl_1\"\n\nencoder_config = SimpleNamespace(**{})\nencoder_config.input_dim=256\nencoder_config.encoder_dim=256\nencoder_config.num_layers=4\nencoder_config.num_attention_heads= 16\nencoder_config.feed_forward_expansion_factor=2\nencoder_config.conv_expansion_factor= 2\nencoder_config.input_dropout_p= 0.0\nencoder_config.feed_forward_dropout_p= 0.0\nencoder_config.attention_dropout_p= 0.0\nencoder_config.conv_dropout_p= 0.0\nencoder_config.conv_kernel_size= 51\n\ncfg.encoder_config = encoder_config\n\ncfg.use_bn= True\ncfg.use_gnn = True \ncfg.cnn_extractor = False\n\n\n# LOSS SETTINGS (for handling sparse labels)\ncfg.use_focal_loss = True  # Start with weighted BCE, simpler to debug\ncfg.focal_alpha = 0.25  # Weight for positive/negative samples in focal loss\ncfg.focal_gamma = 2.0  # Focusing parameter for focal loss\ncfg.pos_weight = 10.0  # High weight for positive samples due to extreme sparsity (0.1% positive)\n\n# OPTIMIZATION & SCHEDULE\ncfg.fold = 0\ncfg.epochs = 100\ncfg.eval_epochs = 5\ncfg.lr = 1e-4  # Reduced learning rate for stability with larger windows\ncfg.optimizer = \"AdamW\"\ncfg.weight_decay = 0.05\ncfg.clip_grad = 0.\ncfg.warmup = 5\ncfg.batch_size = 64\ncfg.batch_size_val = 64\ncfg.mixed_precision = True # True\ncfg.pin_memory = False\ncfg.grad_accumulation = 1.\ncfg.num_workers = 4\n\n\n#EVAL\ncfg.calc_metric = True\ncfg.metric = \"metric_1\"\ncfg.save_val_data = True","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T13:06:42.371423Z","iopub.execute_input":"2025-12-18T13:06:42.371689Z","iopub.status.idle":"2025-12-18T13:06:48.584937Z","shell.execute_reply.started":"2025-12-18T13:06:42.371664Z","shell.execute_reply":"2025-12-18T13:06:48.584314Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### PREPROCESS","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport os\nimport pandas as pd\nimport numpy as np\nfrom scipy.interpolate import interp1d\nimport sys\nimport warnings\n\nwarnings.filterwarnings('ignore', message='Mean of empty slice')\n\nDATA_DIR = '/kaggle/input/MABe-mouse-behavior-detection'\nPREPROC_DIR = '/kaggle/working//preprocessed'\nos.makedirs(PREPROC_DIR, exist_ok=True)\nos.makedirs(os.path.join(PREPROC_DIR, 'test'), exist_ok=True)\n\ntest_meta = pd.read_csv(os.path.join(DATA_DIR, '/kaggle/input/MABe-mouse-behavior-detection/test.csv'))\nall_meta = test_meta.copy()  # For this example, only train\n\nMASTER_SKELETON = [\n    'nose',\n    'ear_left',\n    'ear_right',\n    'head_center',\n    'body_center',\n    'tail_base'\n]\nNUM_MASTER_KEYPOINTS = len(MASTER_SKELETON)\nMASTER_SKELETON_MAP = {name: i for i, name in enumerate(MASTER_SKELETON)}\n\nfixed_bodyparts = cfg.tracked_bodyparts\nnum_bodyparts = len(fixed_bodyparts)\nfixed_actions = cfg.set_behavior_classes\nnum_actions = len(fixed_actions)\naction_map = cfg.action_id_map\n\nNO_ACTION_IDX = action_map['no_action']  # GET THE NO_ACTION INDEX\nprint(f\"Fixed actions ({num_actions}): {fixed_actions}\")\nprint(f\"no_action index: {NO_ACTION_IDX}\")\n# Max mice and pairs from cfg\nset_mice = cfg.set_mice\nmax_mice = len(set_mice)  # 4 if 4 mice\nmax_pairs = cfg.max_pairs  # 16 if 4 mice\n\n# Mouse map from cfg\nmouse_id_map = cfg.mouse_id_map  # Assume {1: 0, 2: 1, 3: 2, 4: 3} or similar\n\n# Frame counts\nframe_counts = {'test': {}}\n\ndef process_tracking_video(row, split):\n    lab_id = row['lab_id']\n    video_id = row['video_id']\n    pix_per_cm = row.get('pix_per_cm_approx', 1.0)\n    \n    if pd.isna(pix_per_cm) or pix_per_cm <= 0:\n        print(f\"Warning {video_id}: Invalid pix_per_cm ({pix_per_cm}), defaulting to 1.0\")\n        pix_per_cm = 1.0\n    \n    # Check for no mice from metadata\n    num_mice = sum(1 for i in range(1, max_mice + 1) if pd.notna(row.get(f'mouse{i}_id')))\n    if num_mice == 0:\n        print(f\"Skipping {video_id}: no mice\")\n        return\n    \n    # Load tracking parquet\n    path = os.path.join(DATA_DIR, f'{split}_tracking', lab_id, f'{video_id}.parquet')\n    if not os.path.exists(path):\n        print(f\"Warning: No tracking file for {video_id}\")\n        return\n    df = pd.read_parquet(path)\n    if df.empty:\n        print(f\"Warning: Empty tracking for {video_id}\")\n        return\n    \n    # Handle types\n    df['video_frame'] = df['video_frame'].astype(int)\n    df['mouse_id'] = df['mouse_id'].astype(int)\n    max_frame = df['video_frame'].max()\n    \n    # Pivot to wide\n    df_x = df.pivot(index='video_frame', columns=['mouse_id', 'bodypart'], values='x')\n    df_y = df.pivot(index='video_frame', columns=['mouse_id', 'bodypart'], values='y')\n    \n    # Fill missing frames\n    full_index = pd.Index(range(0, max_frame + 1), name='video_frame')\n    df_x = df_x.reindex(full_index, fill_value=np.nan)\n    df_y = df_y.reindex(full_index, fill_value=np.nan)\n    num_frames = len(full_index)\n    \n    # Map mouse_ids to 0-based using cfg map\n    current_mouse_levels = df_x.columns.levels[0]\n    new_mouse_levels = [mouse_id_map.get(level, -1) for level in current_mouse_levels]\n    if any(l < 0 for l in new_mouse_levels):\n        print(f\"Warning: Unknown mouse_ids in {video_id}\")\n    df_x.columns = df_x.columns.set_levels(new_mouse_levels, level=0)\n    df_y.columns = df_y.columns.set_levels(new_mouse_levels, level=0)\n    \n    # Helper to get bodypart array [frames, mice], reindexed to all mice with NaN fill\n    def get_bp_array(bp, is_x=True, fill_nan=True):\n        df = df_x if is_x else df_y\n        if bp not in df.columns.levels[1]:\n            return np.full((num_frames, max_mice), np.nan, dtype=np.float32)\n        bp_df = df.xs(bp, level=1, axis=1)\n        bp_df = bp_df.reindex(columns=range(max_mice), fill_value=np.nan)\n        return bp_df.values.astype(np.float32)\n    \n    # For master skeleton, create keypoints array\n    keypoints = np.full((num_frames, max_mice, NUM_MASTER_KEYPOINTS, 2), np.nan, dtype=np.float32)\n    \n    # Vectorized mapping\n    # Nose\n    nose_x = get_bp_array('nose', is_x=True)\n    nose_y = get_bp_array('nose', is_x=False)\n    keypoints[:, :, MASTER_SKELETON_MAP['nose'], 0] = nose_x\n    keypoints[:, :, MASTER_SKELETON_MAP['nose'], 1] = nose_y\n    # Special fallback for labs like GroovyShrew\n    if lab_id == 'GroovyShrew':\n        head_x = get_bp_array('head', is_x=True)\n        head_y = get_bp_array('head', is_x=False)\n        keypoints[:, :, MASTER_SKELETON_MAP['nose'], 0] = np.where(np.isnan(nose_x), head_x, nose_x)\n        keypoints[:, :, MASTER_SKELETON_MAP['nose'], 1] = np.where(np.isnan(nose_y), head_y, nose_y)\n    \n    # Ear left/right\n    keypoints[:, :, MASTER_SKELETON_MAP['ear_left'], 0] = get_bp_array('ear_left', is_x=True)\n    keypoints[:, :, MASTER_SKELETON_MAP['ear_left'], 1] = get_bp_array('ear_left', is_x=False)\n    keypoints[:, :, MASTER_SKELETON_MAP['ear_right'], 0] = get_bp_array('ear_right', is_x=True)\n    keypoints[:, :, MASTER_SKELETON_MAP['ear_right'], 1] = get_bp_array('ear_right', is_x=False)\n    \n    # Tail base\n    keypoints[:, :, MASTER_SKELETON_MAP['tail_base'], 0] = get_bp_array('tail_base', is_x=True)\n    keypoints[:, :, MASTER_SKELETON_MAP['tail_base'], 1] = get_bp_array('tail_base', is_x=False)\n    \n    # Head_Center: first head, then neck, then average ear_left/right\n    head_x = get_bp_array('head', is_x=True)\n    head_y = get_bp_array('head', is_x=False)\n    neck_x = get_bp_array('neck', is_x=True)\n    neck_y = get_bp_array('neck', is_x=False)\n    ear_left_x = keypoints[:, :, MASTER_SKELETON_MAP['ear_left'], 0]\n    ear_left_y = keypoints[:, :, MASTER_SKELETON_MAP['ear_left'], 1]\n    ear_right_x = keypoints[:, :, MASTER_SKELETON_MAP['ear_right'], 0]\n    ear_right_y = keypoints[:, :, MASTER_SKELETON_MAP['ear_right'], 1]\n    \n    avg_ear_x = np.nanmean(np.stack([ear_left_x, ear_right_x], axis=-1), axis=-1)\n    avg_ear_y = np.nanmean(np.stack([ear_left_y, ear_right_y], axis=-1), axis=-1)\n    \n    head_center_x = np.where(~np.isnan(head_x), head_x,\n                    np.where(~np.isnan(neck_x), neck_x, avg_ear_x))\n    head_center_y = np.where(~np.isnan(head_y), head_y,\n                    np.where(~np.isnan(neck_y), neck_y, avg_ear_y))\n    keypoints[:, :, MASTER_SKELETON_MAP['head_center'], 0] = head_center_x\n    keypoints[:, :, MASTER_SKELETON_MAP['head_center'], 1] = head_center_y\n    \n    # Body_Center: body_center, then average [spine_1, spine_2, hip_left, hip_right, neck], then (head_center + tail_base)/2\n    body_center_x = get_bp_array('body_center', is_x=True)\n    body_center_y = get_bp_array('body_center', is_x=False)\n    \n    spine1_x = get_bp_array('spine_1', is_x=True)\n    spine1_y = get_bp_array('spine_1', is_x=False)\n    spine2_x = get_bp_array('spine_2', is_x=True)\n    spine2_y = get_bp_array('spine_2', is_x=False)\n    hip_left_x = get_bp_array('hip_left', is_x=True)\n    hip_left_y = get_bp_array('hip_left', is_x=False)\n    hip_right_x = get_bp_array('hip_right', is_x=True)\n    hip_right_y = get_bp_array('hip_right', is_x=False)\n    neck_x = get_bp_array('neck', is_x=True)  # Reuse\n    neck_y = get_bp_array('neck', is_x=False)\n    \n    avg_body_x = np.nanmean(np.stack([spine1_x, spine2_x, hip_left_x, hip_right_x, neck_x], axis=-1), axis=-1)\n    avg_body_y = np.nanmean(np.stack([spine1_y, spine2_y, hip_left_y, hip_right_y, neck_y], axis=-1), axis=-1)\n    \n    tail_base_x = keypoints[:, :, MASTER_SKELETON_MAP['tail_base'], 0]\n    tail_base_y = keypoints[:, :, MASTER_SKELETON_MAP['tail_base'], 1]\n    head_center_x = keypoints[:, :, MASTER_SKELETON_MAP['head_center'], 0]  # Updated now\n    head_center_y = keypoints[:, :, MASTER_SKELETON_MAP['head_center'], 1]\n    fallback_x = (head_center_x + tail_base_x) / 2\n    fallback_y = (head_center_y + tail_base_y) / 2\n    fallback_x[np.isnan(fallback_x)] = np.nan  # If either missing\n    fallback_y[np.isnan(fallback_y)] = np.nan\n    \n    final_body_x = np.where(~np.isnan(body_center_x), body_center_x,\n                   np.where(~np.isnan(avg_body_x), avg_body_x, fallback_x))\n    final_body_y = np.where(~np.isnan(body_center_y), body_center_y,\n                   np.where(~np.isnan(avg_body_y), avg_body_y, fallback_y))\n    keypoints[:, :, MASTER_SKELETON_MAP['body_center'], 0] = final_body_x\n    keypoints[:, :, MASTER_SKELETON_MAP['body_center'], 1] = final_body_y\n    \n    # Create validity mask before imputation (1: valid/original non-NaN, 0: invalid/NaN)\n    mask = (~np.isnan(keypoints[..., 0])).astype(np.uint8)  # [frames, mice, bp] since x/y same\n    \n    keypoints /= pix_per_cm\n    \n    # print(f\"Processed {video_id}: {num_frames} frames, feature_dim {keypoints.shape[1]}\")\n    return keypoints.astype(np.float32), num_frames","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T13:06:48.585683Z","iopub.execute_input":"2025-12-18T13:06:48.586018Z","iopub.status.idle":"2025-12-18T13:06:48.926729Z","shell.execute_reply.started":"2025-12-18T13:06:48.585997Z","shell.execute_reply":"2025-12-18T13:06:48.925904Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### MODEL OLD\n","metadata":{}},{"cell_type":"code","source":"%%writefile model_1.py\n\nimport torch\nfrom torch.nn import functional as F\nfrom torch import nn\nfrom typing import Tuple, Union, Optional\nfrom torch import Tensor\nimport math\nimport numpy as np\n\n\ndef count_parameters(model):\n    \"\"\"Count the number of trainable parameters in model\"\"\"\n    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n\nclass Swish(nn.Module):\n    \"\"\"Swish activation function\"\"\"\n    def __init__(self) -> None:\n        super(Swish, self).__init__()\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return inputs * inputs.sigmoid()\n\nclass GLU(nn.Module):\n    \"\"\"Gated Linear Unit activation\"\"\"\n    def __init__(self, dim: int) -> None:\n        super(GLU, self).__init__()\n        self.dim = dim\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        outputs, gate = inputs.chunk(2, dim=self.dim)\n        return outputs * gate.sigmoid()\n\nclass FeedForwardModule(nn.Module):\n    \"\"\"\n    Feed Forward Module with pre-norm residual units\n    \"\"\"\n    def __init__(\n        self,\n        encoder_dim: int = 512,\n        expansion_factor: int = 4,\n        dropout_p: float = 0.0,\n    ) -> None:\n        super(FeedForwardModule, self).__init__()\n\n        self.ffn1 = nn.Linear(encoder_dim, encoder_dim * expansion_factor, bias=True)\n        self.act = Swish()\n        self.do1 = nn.Dropout(p=dropout_p)\n        self.ffn2 = nn.Linear(encoder_dim * expansion_factor, encoder_dim, bias=True)\n        self.do2 = nn.Dropout(p=dropout_p)\n\n    def forward(self, x):\n        x = self.ffn1(x)\n        x = self.act(x)\n        x = self.do1(x)\n        x = self.ffn2(x)\n        x = self.do2(x)\n        return x\n\nclass RelPositionalEncoding(nn.Module):\n    \"\"\"\n    Relative positional encoding module for handling variable sequence lengths\n    \"\"\"\n    def __init__(self, d_model: int = 512, max_len: int = 5000) -> None:\n        super(RelPositionalEncoding, self).__init__()\n        self.d_model = d_model\n        self.pe = None\n        self.extend_pe(torch.tensor(0.0).expand(1, max_len))\n\n    def extend_pe(self, x):\n        if self.pe is not None:\n            if self.pe.size(1) >= x.size(1) * 2 - 1:\n                if self.pe.dtype != x.dtype or self.pe.device != x.device:\n                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)\n                return\n\n        pe_positive = torch.zeros(x.size(1), self.d_model)\n        pe_negative = torch.zeros(x.size(1), self.d_model)\n        position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)\n        div_term = torch.exp(\n            torch.arange(0, self.d_model, 2, dtype=torch.float32) * -(math.log(10000.0) / self.d_model)\n        )\n        pe_positive[:, 0::2] = torch.sin(position * div_term)\n        pe_positive[:, 1::2] = torch.cos(position * div_term)\n        pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)\n        pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)\n\n        pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)\n        pe_negative = pe_negative[1:].unsqueeze(0)\n        pe = torch.cat([pe_positive, pe_negative], dim=1)\n        self.pe = pe.to(device=x.device, dtype=x.dtype)\n\n    def forward(self, x: torch.Tensor):\n        self.extend_pe(x)\n        pos_emb = self.pe[\n            :,\n            self.pe.size(1) // 2 - x.size(1) + 1 : self.pe.size(1) // 2 + x.size(1),\n        ]\n        return pos_emb\n\nclass RelativeMultiHeadAttention(nn.Module):\n    \"\"\"\n    Multi-head attention with relative positional encoding\n    \"\"\"\n    def __init__(\n        self,\n        d_model: int = 512,\n        num_heads: int = 16,\n        dropout_p: float = 0.0,\n    ):\n        super(RelativeMultiHeadAttention, self).__init__()\n        assert d_model % num_heads == 0, \"d_model % num_heads should be zero.\"\n        self.d_model = d_model\n        self.d_head = int(d_model / num_heads)\n        self.num_heads = num_heads\n        self.sqrt_dim = math.sqrt(self.d_head)\n\n        self.query_proj = nn.Linear(d_model, d_model)\n        self.key_proj = nn.Linear(d_model, d_model)\n        self.value_proj = nn.Linear(d_model, d_model)\n        self.pos_proj = nn.Linear(d_model, d_model, bias=False)\n\n        self.dropout = nn.Dropout(p=dropout_p)\n        self.u_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))\n        self.v_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))\n        torch.nn.init.xavier_uniform_(self.u_bias)\n        torch.nn.init.xavier_uniform_(self.v_bias)\n\n        self.out_proj = nn.Linear(d_model, d_model)\n\n    def forward(\n        self,\n        query: Tensor,\n        key: Tensor,\n        value: Tensor,\n        pos_embedding: Tensor,\n        mask: Optional[Tensor] = None,\n    ) -> Tensor:\n        batch_size = value.size(0)\n\n        query = self.query_proj(query).view(batch_size, -1, self.num_heads, self.d_head)\n        key = self.key_proj(key).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)\n        value = self.value_proj(value).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)\n        pos_embedding = self.pos_proj(pos_embedding).view(batch_size, -1, self.num_heads, self.d_head)\n\n        content_score = torch.matmul((query + self.u_bias).transpose(1, 2), key.transpose(2, 3))\n        pos_score = torch.matmul((query + self.v_bias).transpose(1, 2), pos_embedding.permute(0, 2, 3, 1))\n        pos_score = self._relative_shift(pos_score)\n\n        score = (content_score + pos_score) / self.sqrt_dim\n\n        if mask is not None:\n            mask = mask.unsqueeze(1)\n            score.masked_fill_(mask, -1e9) if score.dtype == torch.float32 else score.masked_fill_(mask, -1e4)\n\n        attn = F.softmax(score, -1)\n        attn = self.dropout(attn)\n\n        context = torch.matmul(attn, value).transpose(1, 2)\n        context = context.contiguous().view(batch_size, -1, self.d_model)\n\n        return self.out_proj(context)\n\n    def _relative_shift(self, pos_score: Tensor) -> Tensor:\n        batch_size, num_heads, seq_length1, seq_length2 = pos_score.size()\n        zeros = pos_score.new_zeros(batch_size, num_heads, seq_length1, 1)\n        padded_pos_score = torch.cat([zeros, pos_score], dim=-1)\n\n        padded_pos_score = padded_pos_score.view(batch_size, num_heads, seq_length2 + 1, seq_length1)\n        pos_score = padded_pos_score[:, :, 1:].view_as(pos_score)[:, :, :, : seq_length2 // 2 + 1]\n\n        return pos_score\n\nclass MultiHeadedSelfAttentionModule(nn.Module):\n    \"\"\"\n    Self-attention module with relative positional encoding\n    \"\"\"\n    def __init__(self, d_model: int, num_heads: int, dropout_p: float = 0.0):\n        super(MultiHeadedSelfAttentionModule, self).__init__()\n        self.positional_encoding = RelPositionalEncoding(d_model)\n        self.attention = RelativeMultiHeadAttention(d_model, num_heads, dropout_p)\n        self.dropout = nn.Dropout(p=dropout_p)\n\n    def forward(self, inputs: Tensor, mask: Optional[Tensor] = None):\n        batch_size = inputs.size(0)\n        pos_embedding = self.positional_encoding(inputs)\n        pos_embedding = pos_embedding.repeat(batch_size, 1, 1)\n\n        outputs = self.attention(inputs, inputs, inputs, pos_embedding=pos_embedding, mask=mask)\n        return self.dropout(outputs)\n\nclass DepthwiseConv1d(nn.Module):\n    \"\"\"Depthwise 1D convolution\"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        kernel_size: int,\n        stride: int = 1,\n        padding: int = 0,\n        bias: bool = False,\n    ) -> None:\n        super(DepthwiseConv1d, self).__init__()\n        assert out_channels % in_channels == 0, \"out_channels should be constant multiple of in_channels\"\n        self.conv = nn.Conv1d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=kernel_size,\n            groups=in_channels,\n            stride=stride,\n            padding=padding,\n            bias=bias,\n        )\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return self.conv(inputs)\n\nclass PointwiseConv1d(nn.Module):\n    \"\"\"Pointwise 1D convolution (kernel size = 1)\"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        stride: int = 1,\n        padding: int = 0,\n        bias: bool = True,\n    ) -> None:\n        super(PointwiseConv1d, self).__init__()\n        self.conv = nn.Conv1d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=1,\n            stride=stride,\n            padding=padding,\n            bias=bias,\n        )\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return self.conv(inputs)\n\nclass ConvModule(nn.Module):\n    \"\"\"\n    Convolution module with pointwise conv -> GLU -> depthwise conv -> normalization -> activation\n    \"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        kernel_size: int = 31,\n        expansion_factor: int = 2,\n        dropout_p: float = 0.0,\n        use_bn: bool = True,\n    ) -> None:\n        super(ConvModule, self).__init__()\n        assert (kernel_size - 1) % 2 == 0, \"kernel_size should be a odd number for 'SAME' padding\"\n        assert expansion_factor == 2, \"Currently, Only Supports expansion_factor 2\"\n\n        self.pw_conv_1 = PointwiseConv1d(in_channels, in_channels * expansion_factor, stride=1, padding=0, bias=True)\n        self.act1 = GLU(dim=1)\n        self.dw_conv = DepthwiseConv1d(in_channels, in_channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2)\n        self.bn = nn.BatchNorm1d(in_channels)\n        self.inorm = nn.InstanceNorm1d(in_channels, affine=True)\n        self.act2 = Swish()\n        self.pw_conv_2 = PointwiseConv1d(in_channels, in_channels, stride=1, padding=0, bias=True)\n        self.do = nn.Dropout(p=dropout_p)\n        self.use_bn = use_bn\n\n    def forward(self, x, mask_pad):\n        # Transpose for conv operations [B, T, C]\n        x = x.transpose(1, 2)\n        if mask_pad.size(2) > 0:  # time > 0\n            x = x.masked_fill(~mask_pad, 0.0)\n\n        x = self.pw_conv_1(x)\n        x = self.act1(x)\n        x = self.dw_conv(x)\n\n        if self.use_bn:\n            # Apply batch norm only to non-padded positions\n            x_bn = x.permute(0,2,1).reshape(-1, x.shape[1])\n            mask_bn = mask_pad.view(-1)\n            x_bn[mask_bn] = self.bn(x_bn[mask_bn])\n            x = x_bn.view(x.permute(0,2,1).shape).permute(0,2,1)\n        else:    \n            x = self.inorm(x)\n\n        x = self.act2(x)\n        x = self.pw_conv_2(x)\n        x = self.do(x)\n\n        # Mask batch padding again\n        if mask_pad.size(2) > 0:  # time > 0\n            x = x.masked_fill(~mask_pad, 0.0)\n        x = x.transpose(1, 2)\n        return x\n\ndef make_scale(encoder_dim):\n    \"\"\"Create learnable scale and bias parameters\"\"\"\n    scale = torch.nn.Parameter(torch.tensor([1.] * encoder_dim)[None, None, :])\n    bias = torch.nn.Parameter(torch.tensor([0.] * encoder_dim)[None, None, :])\n    return scale, bias\n\nclass SqueezeformerBlock(nn.Module):\n    \"\"\"\n    Squeezeformer block: MHSA -> FFN -> Conv -> FFN with residual connections\n    \"\"\"\n    def __init__(\n        self,\n        encoder_dim: int = 512,\n        num_attention_heads: int = 8,\n        feed_forward_expansion_factor: int = 4,\n        conv_expansion_factor: int = 2,\n        feed_forward_dropout_p: float = 0.0,\n        attention_dropout_p: float = 0.1,\n        conv_dropout_p: float = 0.1,\n        conv_kernel_size: int = 31,\n        use_bn: bool = True,\n    ):\n        super(SqueezeformerBlock, self).__init__()\n\n        self.scale_mhsa, self.bias_mhsa = make_scale(encoder_dim)\n        self.scale_ff_mhsa, self.bias_ff_mhsa = make_scale(encoder_dim)\n        self.scale_conv, self.bias_conv = make_scale(encoder_dim)\n        self.scale_ff_conv, self.bias_ff_conv = make_scale(encoder_dim)\n\n        self.mhsa = MultiHeadedSelfAttentionModule(\n            d_model=encoder_dim,\n            num_heads=num_attention_heads,\n            dropout_p=attention_dropout_p,\n        )\n        self.ln_mhsa = nn.LayerNorm(encoder_dim)\n        self.ff_mhsa = FeedForwardModule(\n            encoder_dim=encoder_dim,\n            expansion_factor=feed_forward_expansion_factor,\n            dropout_p=feed_forward_dropout_p,\n        )\n        self.ln_ff_mhsa = nn.LayerNorm(encoder_dim)\n        self.conv = ConvModule(\n            in_channels=encoder_dim,\n            kernel_size=conv_kernel_size,\n            expansion_factor=conv_expansion_factor,\n            dropout_p=conv_dropout_p,\n            use_bn=use_bn,\n        )\n        self.ln_conv = nn.LayerNorm(encoder_dim)\n        self.ff_conv = FeedForwardModule(\n            encoder_dim=encoder_dim,\n            expansion_factor=feed_forward_expansion_factor,\n            dropout_p=feed_forward_dropout_p,\n        )\n        self.ln_ff_conv = nn.LayerNorm(encoder_dim)\n\n    def forward(self, x, mask):\n        mask_pad = (mask).long().bool().unsqueeze(1)\n        mask_pad = ~(mask_pad.permute(0, 2, 1) * mask_pad)\n        mask_flat = mask.view(-1).bool()\n        bs, slen, nfeats = x.shape\n\n        # MHSA\n        residual = x\n        x = x * self.scale_mhsa + self.bias_mhsa\n        x = residual + self.mhsa(x, mask_pad)\n\n        # Skip padding for layer norm\n        x_skip = x.view(-1, x.shape[-1])\n        x = x_skip[mask_flat].unsqueeze(0)\n        x = self.ln_mhsa(x)\n\n        # FFN after MHSA\n        residual = x\n        x = x * self.scale_ff_mhsa + self.bias_ff_mhsa\n        x = residual + self.ff_mhsa(x)\n        x = self.ln_ff_mhsa(x)\n\n        # Restore shape\n        x_skip[mask_flat] = x[0]\n        x = x_skip.view(bs, slen, nfeats)\n\n        # Conv\n        residual = x\n        x = x * self.scale_conv + self.bias_conv\n        x = residual + self.conv(x, mask_pad=mask.bool().unsqueeze(1))\n\n        # Skip padding for layer norm\n        x_skip = x.view(-1, x.shape[-1])\n        x = x_skip[mask_flat].unsqueeze(0)\n        x = self.ln_conv(x)\n\n        # FFN after Conv\n        residual = x\n        x = x * self.scale_ff_conv + self.bias_ff_conv\n        x = residual + self.ff_conv(x)\n        x = self.ln_ff_conv(x)\n\n        # Restore shape\n        x_skip[mask_flat] = x[0]\n        x = x_skip.view(bs, slen, nfeats)\n\n        return x\n\nclass SqueezeformerEncoder(nn.Module):\n    \"\"\"\n    Stack of Squeezeformer blocks\n    \"\"\"\n    def __init__(\n        self,\n        input_dim: int = 80,\n        encoder_dim: int = 512,\n        num_layers: int = 16,\n        num_attention_heads: int = 8,\n        feed_forward_expansion_factor: int = 4,\n        conv_expansion_factor: int = 2,\n        input_dropout_p: float = 0.0,\n        feed_forward_dropout_p: float = 0.0,\n        attention_dropout_p: float = 0.0,\n        conv_dropout_p: float = 0.0,\n        conv_kernel_size: int = 31,\n        use_bn: bool = True,\n    ):\n        super(SqueezeformerEncoder, self).__init__()\n        self.num_layers = num_layers\n\n        self.blocks = nn.ModuleList()\n        for idx in range(num_layers):\n            self.blocks.append(\n                SqueezeformerBlock(\n                    encoder_dim=encoder_dim,\n                    num_attention_heads=num_attention_heads,\n                    feed_forward_expansion_factor=feed_forward_expansion_factor,\n                    conv_expansion_factor=conv_expansion_factor,\n                    feed_forward_dropout_p=feed_forward_dropout_p,\n                    attention_dropout_p=attention_dropout_p,\n                    conv_dropout_p=conv_dropout_p,\n                    conv_kernel_size=conv_kernel_size,\n                    use_bn=use_bn,\n                )\n            )\n\n    def forward(self, x: Tensor, mask: Tensor):\n        for idx, block in enumerate(self.blocks):\n            x = block(x, mask)\n        return x\n\nfrom timm.layers.norm_act import BatchNormAct2d\n\nclass FeatureExtractor(nn.Module):\n    def __init__(self,\n                 n_landmarks,\n                 out_dim):\n        super().__init__()   \n\n        self.in_channels = in_channels = (32//2) * n_landmarks\n        self.stem_linear = nn.Linear(in_channels,out_dim,bias=False)\n        self.stem_bn = nn.BatchNorm1d(out_dim, momentum=0.95)\n        self.conv_stem = nn.Conv2d(4, 32, kernel_size=(3, 3), stride=(1, 2), padding=(1, 1), bias=False)\n        self.bn_conv = BatchNormAct2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True,act_layer = nn.SiLU,drop_layer=None)\n        \n    def forward(self, data, mask):\n        xc = data.permute(0,2,1,3)  # B,C,T,F\n        xc = self.conv_stem(xc)\n        xc = self.bn_conv(xc)\n        xc = xc.permute(0,2,3,1)\n        xc = xc.reshape(*data.shape[:2], -1)\n        \n        m = mask.to(torch.bool)  \n        x = self.stem_linear(xc)\n        \n        # Batchnorm without pads\n        bs,slen,nfeat = x.shape\n        x = x.view(-1, nfeat)\n        x_bn = x[mask.view(-1)==1].unsqueeze(0)\n        x_bn = self.stem_bn(x_bn.permute(0,2,1)).permute(0,2,1)\n        x[mask.view(-1)==1] = x_bn[0]\n        x = x.view(bs,slen,nfeat)\n        # Padding mask\n        x = x.masked_fill(~mask.bool().unsqueeze(-1), 0.0)\n        \n        return x\n\nclass MABeFeatureExtractor(nn.Module):\n    def __init__(self, input_dim=708, encoder_dim=144, dropout=0.0):\n        super().__init__()\n\n        self.input_dim = input_dim\n        self.encoder_dim = encoder_dim\n\n        # Project input features to encoder dimension\n        self.input_proj = nn.Linear(input_dim, encoder_dim)\n        self.input_norm = nn.LayerNorm(encoder_dim)\n        self.input_dropout = nn.Dropout(dropout)\n\n        # Optional: Add a small CNN for local temporal patterns\n        self.use_conv = True\n        if self.use_conv:\n            self.conv1 = nn.Conv1d(encoder_dim, encoder_dim, kernel_size=3, padding=1)\n            self.conv2 = nn.Conv1d(encoder_dim, encoder_dim, kernel_size=3, padding=1)\n            self.conv_norm = nn.LayerNorm(encoder_dim)\n            self.conv_dropout = nn.Dropout(dropout)\n\n    def forward(self, x, mask):\n        \"\"\"\n        Args:\n            x: Input features (batch, seq_len, 708)\n            mask: Attention mask (batch, seq_len)\n        Returns:\n            Encoded features (batch, seq_len, encoder_dim)\n        \"\"\"\n        # Project to encoder dimension\n        x = self.input_proj(x)\n        x = self.input_norm(x)\n        x = self.input_dropout(x)\n\n        if self.use_conv:\n            # Apply temporal convolutions\n            x_conv = x.transpose(1, 2)  # (batch, encoder_dim, seq_len)\n            x_conv = F.relu(self.conv1(x_conv))\n            x_conv = F.relu(self.conv2(x_conv))\n            x_conv = x_conv.transpose(1, 2)  # (batch, seq_len, encoder_dim)\n\n            # Residual connection\n            x = x + self.conv_dropout(self.conv_norm(x_conv))\n\n        # Apply mask\n        x = x.masked_fill(~mask.bool().unsqueeze(-1), 0.0)\n\n        return x\n\nimport torch.nn.functional as F\nclass GraphAggregation(nn.Module):\n    def __init__(self, node_dim):\n        super(GraphAggregation, self).__init__()\n        M = 4  # Fixed num_mice\n        self.linear = nn.Linear(node_dim, node_dim)  # Message transformation\n        self.adjacency = nn.Parameter(torch.eye(M))  # Learnable adj matrix with self-loops\n\n    def forward(self, nodes):\n        \"\"\"\n        nodes: [B*T, M, node_dim] batched per-mouse features.\n        Returns aggregated [B*T, M, node_dim].\n        \"\"\"\n        messages = self.linear(nodes)  # [B*T, M, node_dim]\n        adj_batch = self.adjacency.unsqueeze(0).expand(nodes.size(0), -1, -1)  # [B*T, M, M]\n        aggregated = torch.bmm(adj_batch, messages)  # [B*T, M, M] @ [B*T, M, dim] = [B*T, M, dim]\n        return aggregated\n\nclass BehaviorClassificationHead(nn.Module):\n\n    def __init__(self, encoder_dim=144, num_pairs=16, num_actions=39, dropout=0.0):\n        super().__init__()\n\n        self.num_pairs = num_pairs  # Number of mouse pairs (e.g., 4 mice -> 16 directed pairs)\n        self.num_actions = num_actions  # Number of behavior classes\n\n        # Shared layers for all pairs\n        self.shared_proj = nn.Sequential(\n            nn.Linear(encoder_dim, encoder_dim * 2),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(encoder_dim * 2, encoder_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout)\n        )\n\n        # Single head that outputs all pairs x actions\n        # This allows better parameter sharing across pairs\n        self.classifier = nn.Linear(encoder_dim, num_pairs * num_actions)\n\n    def forward(self, x, mask=None):\n        \"\"\"\n        Args:\n            x: Encoder output (batch, seq_len, encoder_dim)\n            mask: Attention mask (batch, seq_len)\n        Returns:\n            Logits (batch, seq_len, num_pairs, num_actions)\n        \"\"\"\n        batch_size, seq_len, _ = x.shape\n\n        # Shared transformation\n        x = self.shared_proj(x)  # (batch, seq_len, encoder_dim)\n\n        # Single classifier for all pairs\n        logits = self.classifier(x)  # (batch, seq_len, num_pairs * num_actions)\n\n        # Reshape to separate pairs and actions\n        logits = logits.view(batch_size, seq_len, self.num_pairs, self.num_actions)\n\n        return logits\n\nclass Net(nn.Module):\n    \"\"\"\n    Squeezeformer model for MABe mouse behavior detection\n    \"\"\"\n    def __init__(self, cfg):\n        super(Net, self).__init__()\n        self.cfg = cfg\n        # Model dimensions\n        self.encoder_dim = cfg.encoder_config.encoder_dim\n        self.num_pairs = 16  # Number of mouse pairs\n        self.num_actions = 39  # Number of behavior classes (including no_action)\n\n        self.gnn = GraphAggregation(self.encoder_dim // 4) if getattr(cfg, 'use_gnn', False) else None\n\n        # Feature extractor\n        if self.cfg.cnn_extractor:\n            self.feature_extractor = FeatureExtractor(\n                n_landmarks=cfg.per_mouse_feature_dim // 4,\n                out_dim=cfg.encoder_config.encoder_dim)\n        else:\n            self.feature_extractor = MABeFeatureExtractor(\n                input_dim=cfg.feature_dim,\n                encoder_dim=self.encoder_dim,\n                dropout=cfg.encoder_config.input_dropout_p\n            )\n\n        # Squeezeformer encoder\n        self.encoder = SqueezeformerEncoder(\n            input_dim=self.encoder_dim,\n            encoder_dim=self.encoder_dim,\n            num_layers=cfg.encoder_config.num_layers,\n            num_attention_heads=cfg.encoder_config.num_attention_heads,\n            feed_forward_expansion_factor=cfg.encoder_config.feed_forward_expansion_factor,\n            conv_expansion_factor=cfg.encoder_config.conv_expansion_factor,\n            input_dropout_p=cfg.encoder_config.input_dropout_p,\n            feed_forward_dropout_p=cfg.encoder_config.feed_forward_dropout_p,\n            attention_dropout_p=cfg.encoder_config.attention_dropout_p,\n            conv_dropout_p=cfg.encoder_config.conv_dropout_p,\n            conv_kernel_size=cfg.encoder_config.conv_kernel_size,\n            use_bn = cfg.use_bn\n        )\n\n        # Classification head\n        self.classifier = BehaviorClassificationHead(\n            encoder_dim=self.encoder_dim,\n            num_pairs=self.num_pairs,\n            num_actions=self.num_actions,\n            dropout=cfg.encoder_config.feed_forward_dropout_p\n        )\n\n        if hasattr(cfg, 'class_weights') and cfg.class_weights is not None:\n            class_weights = torch.tensor(cfg.class_weights)\n        else:\n            pos_weight = cfg.pos_weight if hasattr(cfg, 'pos_weight') else 50.0 \n            class_weights = torch.full((self.num_actions,), pos_weight) # High weight for all\n            no_action_idx = cfg.action_id_map.get('no_action', -1)\n            if no_action_idx != -1:\n                class_weights[no_action_idx] = 1.0 # Low weight for no_action\n            else:\n                # Fallback if cfg is wrong, assume last class is no_action\n                class_weights[-1] = 1.0\n\n        self.register_buffer('class_weights', class_weights)\n\n        # Use CrossEntropyLoss for multi-class classification\n        self.loss_fn = nn.CrossEntropyLoss(weight=self.class_weights, reduction='none')\n        self.unweighted_ce_fn = nn.CrossEntropyLoss(reduction='none')\n\n        # Optional: Multi-class focal loss\n        self.use_focal_loss = cfg.use_focal_loss if hasattr(cfg, 'use_focal_loss') else False\n        self.focal_gamma = cfg.focal_gamma if hasattr(cfg, 'focal_gamma') else 2.0\n\n        # Training settings\n        self.return_logits = cfg.return_logits if hasattr(cfg, 'return_logits') else False\n\n        print(f'Model initialized with {count_parameters(self):,} trainable parameters')\n        # print(f'Loss: {\"Multi-class Focal\" if self.use_focal_loss else \"Weighted CrossEntropy\"}')\n        # print(f'Class weights: {self.class_weights.cpu().numpy()}')\n\n    def forward(self, batch):\n        mask = batch['input_mask']  # (batch, seq_len)\n        if self.cfg.cnn_extractor:\n            x = batch['input_mice']\n            B, T, M, _ = x.shape\n            x = self.feature_extractor(x, mask)  # (batch, seq_len, encoder_dim)\n        else:\n            x = batch['input']  # (batch, seq_len, 708)\n            B, T, _ = x.shape\n            x = self.feature_extractor(x, mask)  # (batch, seq_len, encoder_dim)\n\n\n        # Encode\n        x = self.encoder(x, mask)  # (batch, seq_len, encoder_dim)\n\n        if self.gnn is not None:\n            B, T, D = x.shape\n            M = 4  # num_mice\n            node_dim = D // M\n            # Reshape to per-mouse nodes [B*T, M, node_dim]\n            x_resh = x.view(B * T, M, node_dim)\n            # Apply batched GNN\n            x_gnn = self.gnn(x_resh)  # [B*T, M, node_dim]\n            x_gnn = x_gnn.view(B, T, D)  # Back to original shape\n            x = x + x_gnn  # Residual add for fusion\n\n        # Classify\n        logits = self.classifier(x, mask)  # (batch, seq_len, num_pairs, num_actions)\n\n        # APPLY BEHAVIOR MASK if provided (for evaluation)\n        if 'behavior_mask' in batch and batch['behavior_mask'] is not None:\n            behavior_mask = batch['behavior_mask']  # (batch, num_pairs, num_actions)\n            # Expand to match logits shape\n            behavior_mask_expanded = behavior_mask.unsqueeze(1).expand_as(logits)\n            # For multi-class, we need to ensure at least one action is valid per pair\n            # Set masked positions to very negative value\n            logits = torch.where(behavior_mask_expanded.bool(), logits,\n                                torch.tensor(-1e10, dtype=logits.dtype, device=logits.device) if logits.dtype == torch.float32\n                                else torch.tensor(-1e4, dtype=logits.dtype, device=logits.device))\n\n        output = {}\n\n        # Calculate loss if labels provided\n        if 'labels' in batch and batch['labels'] is not None:\n            labels = batch['labels']\n\n            # Convert one-hot to class indices if needed\n            if labels.dim() == 4:  # One-hot encoded (batch, seq_len, num_pairs, num_actions)\n                # Convert to class indices\n                labels = torch.argmax(labels, dim=-1)  # (batch, seq_len, num_pairs)\n\n            # Reshape for loss calculation\n            batch_size, seq_len, num_pairs, num_actions = logits.shape\n            logits_flat = logits.reshape(-1, num_actions)  # (batch*seq*pairs, num_actions)\n            labels_flat = labels.reshape(-1)  # (batch*seq*pairs,)\n\n            # Calculate loss\n            if self.use_focal_loss:\n                # Multi-class focal loss\n                unweighted_ce = self.unweighted_ce_fn(logits_flat, labels_flat)\n                pt = torch.exp(-unweighted_ce)\n                alpha = self.class_weights[labels_flat]  # Gather alpha for true classes\n                focal_loss = alpha * (1 - pt) ** self.focal_gamma * unweighted_ce\n                loss = focal_loss\n            else:\n                # Standard weighted cross-entropy\n                loss = self.loss_fn(logits_flat, labels_flat)\n\n            # Reshape loss back\n            loss = loss.view(batch_size, seq_len, num_pairs)\n\n            mask_expanded = mask.unsqueeze(-1).expand_as(loss)\n            loss = loss * mask_expanded\n\n            # Average over valid positions\n            valid_positions = mask_expanded.sum()\n            if valid_positions > 0:\n                loss = loss.sum() / valid_positions\n            else:\n                loss = loss.sum()\n            output['loss'] = loss\n\n        probs = torch.softmax(logits, dim=-1)  # Multi-class probabilities\n        output['predictions'] = probs\n        output['logits'] = logits\n\n        # Store additional info for evaluation\n        if 'video_id' in batch:\n            output['video_id'] = batch['video_id']\n        if 'start_frame' in batch:\n            output['start_frame'] = batch['start_frame']\n\n        return output","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T13:06:48.927795Z","iopub.execute_input":"2025-12-18T13:06:48.928327Z","iopub.status.idle":"2025-12-18T13:06:48.943831Z","shell.execute_reply.started":"2025-12-18T13:06:48.928295Z","shell.execute_reply":"2025-12-18T13:06:48.943056Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### MODEL CONVTRANSFORMER","metadata":{}},{"cell_type":"code","source":"%%writefile model_gnn.py\n\nimport torch\nfrom torch.nn import functional as F\nfrom torch import nn\nfrom typing import Tuple, Union, Optional\nfrom torch import Tensor\nimport math\nimport numpy as np\n\n\ndef count_parameters(model):\n    \"\"\"Count the number of trainable parameters in model\"\"\"\n    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n\nclass Swish(nn.Module):\n    \"\"\"Swish activation function\"\"\"\n    def __init__(self) -> None:\n        super(Swish, self).__init__()\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return inputs * inputs.sigmoid()\n\nclass GLU(nn.Module):\n    \"\"\"Gated Linear Unit activation\"\"\"\n    def __init__(self, dim: int) -> None:\n        super(GLU, self).__init__()\n        self.dim = dim\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        outputs, gate = inputs.chunk(2, dim=self.dim)\n        return outputs * gate.sigmoid()\n\nclass FeedForwardModule(nn.Module):\n    \"\"\"\n    Feed Forward Module with pre-norm residual units\n    \"\"\"\n    def __init__(\n        self,\n        encoder_dim: int = 512,\n        expansion_factor: int = 4,\n        dropout_p: float = 0.0,\n    ) -> None:\n        super(FeedForwardModule, self).__init__()\n\n        self.ffn1 = nn.Linear(encoder_dim, encoder_dim * expansion_factor, bias=True)\n        self.act = Swish()\n        self.do1 = nn.Dropout(p=dropout_p)\n        self.ffn2 = nn.Linear(encoder_dim * expansion_factor, encoder_dim, bias=True)\n        self.do2 = nn.Dropout(p=dropout_p)\n\n    def forward(self, x):\n        x = self.ffn1(x)\n        x = self.act(x)\n        x = self.do1(x)\n        x = self.ffn2(x)\n        x = self.do2(x)\n        return x\n\nclass RelPositionalEncoding(nn.Module):\n    \"\"\"\n    Relative positional encoding module for handling variable sequence lengths\n    \"\"\"\n    def __init__(self, d_model: int = 512, max_len: int = 5000) -> None:\n        super(RelPositionalEncoding, self).__init__()\n        self.d_model = d_model\n        self.pe = None\n        self.extend_pe(torch.tensor(0.0).expand(1, max_len))\n\n    def extend_pe(self, x):\n        if self.pe is not None:\n            if self.pe.size(1) >= x.size(1) * 2 - 1:\n                if self.pe.dtype != x.dtype or self.pe.device != x.device:\n                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)\n                return\n\n        pe_positive = torch.zeros(x.size(1), self.d_model)\n        pe_negative = torch.zeros(x.size(1), self.d_model)\n        position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)\n        div_term = torch.exp(\n            torch.arange(0, self.d_model, 2, dtype=torch.float32) * -(math.log(10000.0) / self.d_model)\n        )\n        pe_positive[:, 0::2] = torch.sin(position * div_term)\n        pe_positive[:, 1::2] = torch.cos(position * div_term)\n        pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)\n        pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)\n\n        pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)\n        pe_negative = pe_negative[1:].unsqueeze(0)\n        pe = torch.cat([pe_positive, pe_negative], dim=1)\n        self.pe = pe.to(device=x.device, dtype=x.dtype)\n\n    def forward(self, x: torch.Tensor):\n        self.extend_pe(x)\n        pos_emb = self.pe[\n            :,\n            self.pe.size(1) // 2 - x.size(1) + 1 : self.pe.size(1) // 2 + x.size(1),\n        ]\n        return pos_emb\n\nclass RelativeMultiHeadAttention(nn.Module):\n    \"\"\"\n    Multi-head attention with relative positional encoding\n    \"\"\"\n    def __init__(\n        self,\n        d_model: int = 512,\n        num_heads: int = 16,\n        dropout_p: float = 0.0,\n    ):\n        super(RelativeMultiHeadAttention, self).__init__()\n        assert d_model % num_heads == 0, \"d_model % num_heads should be zero.\"\n        self.d_model = d_model\n        self.d_head = int(d_model / num_heads)\n        self.num_heads = num_heads\n        self.sqrt_dim = math.sqrt(self.d_head)\n\n        self.query_proj = nn.Linear(d_model, d_model)\n        self.key_proj = nn.Linear(d_model, d_model)\n        self.value_proj = nn.Linear(d_model, d_model)\n        self.pos_proj = nn.Linear(d_model, d_model, bias=False)\n\n        self.dropout = nn.Dropout(p=dropout_p)\n        self.u_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))\n        self.v_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))\n        torch.nn.init.xavier_uniform_(self.u_bias)\n        torch.nn.init.xavier_uniform_(self.v_bias)\n\n        self.out_proj = nn.Linear(d_model, d_model)\n\n    def forward(\n        self,\n        query: Tensor,\n        key: Tensor,\n        value: Tensor,\n        pos_embedding: Tensor,\n        mask: Optional[Tensor] = None,\n    ) -> Tensor:\n        batch_size = value.size(0)\n\n        query = self.query_proj(query).view(batch_size, -1, self.num_heads, self.d_head)\n        key = self.key_proj(key).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)\n        value = self.value_proj(value).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)\n        pos_embedding = self.pos_proj(pos_embedding).view(batch_size, -1, self.num_heads, self.d_head)\n\n        content_score = torch.matmul((query + self.u_bias).transpose(1, 2), key.transpose(2, 3))\n        pos_score = torch.matmul((query + self.v_bias).transpose(1, 2), pos_embedding.permute(0, 2, 3, 1))\n        pos_score = self._relative_shift(pos_score)\n\n        score = (content_score + pos_score) / self.sqrt_dim\n\n        if mask is not None:\n            mask = mask.unsqueeze(1)\n            score.masked_fill_(mask, -1e9) if score.dtype == torch.float32 else score.masked_fill_(mask, -1e4)\n\n        attn = F.softmax(score, -1)\n        attn = self.dropout(attn)\n\n        context = torch.matmul(attn, value).transpose(1, 2)\n        context = context.contiguous().view(batch_size, -1, self.d_model)\n\n        return self.out_proj(context)\n\n    def _relative_shift(self, pos_score: Tensor) -> Tensor:\n        batch_size, num_heads, seq_length1, seq_length2 = pos_score.size()\n        zeros = pos_score.new_zeros(batch_size, num_heads, seq_length1, 1)\n        padded_pos_score = torch.cat([zeros, pos_score], dim=-1)\n\n        padded_pos_score = padded_pos_score.view(batch_size, num_heads, seq_length2 + 1, seq_length1)\n        pos_score = padded_pos_score[:, :, 1:].view_as(pos_score)[:, :, :, : seq_length2 // 2 + 1]\n\n        return pos_score\n\nclass MultiHeadedSelfAttentionModule(nn.Module):\n    \"\"\"\n    Self-attention module with relative positional encoding\n    \"\"\"\n    def __init__(self, d_model: int, num_heads: int, dropout_p: float = 0.0):\n        super(MultiHeadedSelfAttentionModule, self).__init__()\n        self.positional_encoding = RelPositionalEncoding(d_model)\n        self.attention = RelativeMultiHeadAttention(d_model, num_heads, dropout_p)\n        self.dropout = nn.Dropout(p=dropout_p)\n\n    def forward(self, inputs: Tensor, mask: Optional[Tensor] = None):\n        batch_size = inputs.size(0)\n        pos_embedding = self.positional_encoding(inputs)\n        pos_embedding = pos_embedding.repeat(batch_size, 1, 1)\n\n        outputs = self.attention(inputs, inputs, inputs, pos_embedding=pos_embedding, mask=mask)\n        return self.dropout(outputs)\n\nclass DepthwiseConv1d(nn.Module):\n    \"\"\"Depthwise 1D convolution\"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        kernel_size: int,\n        stride: int = 1,\n        padding: int = 0,\n        bias: bool = False,\n    ) -> None:\n        super(DepthwiseConv1d, self).__init__()\n        assert out_channels % in_channels == 0, \"out_channels should be constant multiple of in_channels\"\n        self.conv = nn.Conv1d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=kernel_size,\n            groups=in_channels,\n            stride=stride,\n            padding=padding,\n            bias=bias,\n        )\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return self.conv(inputs)\n\nclass PointwiseConv1d(nn.Module):\n    \"\"\"Pointwise 1D convolution (kernel size = 1)\"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        stride: int = 1,\n        padding: int = 0,\n        bias: bool = True,\n    ) -> None:\n        super(PointwiseConv1d, self).__init__()\n        self.conv = nn.Conv1d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=1,\n            stride=stride,\n            padding=padding,\n            bias=bias,\n        )\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return self.conv(inputs)\n\nclass ConvModule(nn.Module):\n    \"\"\"\n    Convolution module with pointwise conv -> GLU -> depthwise conv -> normalization -> activation\n    \"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        kernel_size: int = 31,\n        expansion_factor: int = 2,\n        dropout_p: float = 0.0,\n        use_bn: bool = True,\n    ) -> None:\n        super(ConvModule, self).__init__()\n        assert (kernel_size - 1) % 2 == 0, \"kernel_size should be a odd number for 'SAME' padding\"\n        assert expansion_factor == 2, \"Currently, Only Supports expansion_factor 2\"\n\n        self.pw_conv_1 = PointwiseConv1d(in_channels, in_channels * expansion_factor, stride=1, padding=0, bias=True)\n        self.act1 = GLU(dim=1)\n        self.dw_conv = DepthwiseConv1d(in_channels, in_channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2)\n        self.bn = nn.BatchNorm1d(in_channels)\n        self.inorm = nn.InstanceNorm1d(in_channels, affine=True)\n        self.act2 = Swish()\n        self.pw_conv_2 = PointwiseConv1d(in_channels, in_channels, stride=1, padding=0, bias=True)\n        self.do = nn.Dropout(p=dropout_p)\n        self.use_bn = use_bn\n\n    def forward(self, x, mask_pad):\n        # Transpose for conv operations [B, T, C]\n        x = x.transpose(1, 2)\n        if mask_pad.size(2) > 0:  # time > 0\n            x = x.masked_fill(~mask_pad, 0.0)\n\n        x = self.pw_conv_1(x)\n        x = self.act1(x)\n        x = self.dw_conv(x)\n\n        if self.use_bn:\n            # Apply batch norm only to non-padded positions\n            x_bn = x.permute(0,2,1).reshape(-1, x.shape[1])\n            mask_bn = mask_pad.view(-1)\n            x_bn[mask_bn] = self.bn(x_bn[mask_bn])\n            x = x_bn.view(x.permute(0,2,1).shape).permute(0,2,1)\n        else:    \n            x = self.inorm(x)\n\n        x = self.act2(x)\n        x = self.pw_conv_2(x)\n        x = self.do(x)\n\n        # Mask batch padding again\n        if mask_pad.size(2) > 0:  # time > 0\n            x = x.masked_fill(~mask_pad, 0.0)\n        x = x.transpose(1, 2)\n        return x\n\ndef make_scale(encoder_dim):\n    \"\"\"Create learnable scale and bias parameters\"\"\"\n    scale = torch.nn.Parameter(torch.tensor([1.] * encoder_dim)[None, None, :])\n    bias = torch.nn.Parameter(torch.tensor([0.] * encoder_dim)[None, None, :])\n    return scale, bias\n\nclass SqueezeformerBlock(nn.Module):\n    \"\"\"\n    Squeezeformer block: MHSA -> FFN -> Conv -> FFN with residual connections\n    \"\"\"\n    def __init__(\n        self,\n        encoder_dim: int = 512,\n        num_attention_heads: int = 8,\n        feed_forward_expansion_factor: int = 4,\n        conv_expansion_factor: int = 2,\n        feed_forward_dropout_p: float = 0.0,\n        attention_dropout_p: float = 0.1,\n        conv_dropout_p: float = 0.1,\n        conv_kernel_size: int = 31,\n        use_bn: bool = True,\n    ):\n        super(SqueezeformerBlock, self).__init__()\n\n        self.scale_mhsa, self.bias_mhsa = make_scale(encoder_dim)\n        self.scale_ff_mhsa, self.bias_ff_mhsa = make_scale(encoder_dim)\n        self.scale_conv, self.bias_conv = make_scale(encoder_dim)\n        self.scale_ff_conv, self.bias_ff_conv = make_scale(encoder_dim)\n\n        self.mhsa = MultiHeadedSelfAttentionModule(\n            d_model=encoder_dim,\n            num_heads=num_attention_heads,\n            dropout_p=attention_dropout_p,\n        )\n        self.ln_mhsa = nn.LayerNorm(encoder_dim)\n        self.ff_mhsa = FeedForwardModule(\n            encoder_dim=encoder_dim,\n            expansion_factor=feed_forward_expansion_factor,\n            dropout_p=feed_forward_dropout_p,\n        )\n        self.ln_ff_mhsa = nn.LayerNorm(encoder_dim)\n        self.conv = ConvModule(\n            in_channels=encoder_dim,\n            kernel_size=conv_kernel_size,\n            expansion_factor=conv_expansion_factor,\n            dropout_p=conv_dropout_p,\n            use_bn=use_bn,\n        )\n        self.ln_conv = nn.LayerNorm(encoder_dim)\n        self.ff_conv = FeedForwardModule(\n            encoder_dim=encoder_dim,\n            expansion_factor=feed_forward_expansion_factor,\n            dropout_p=feed_forward_dropout_p,\n        )\n        self.ln_ff_conv = nn.LayerNorm(encoder_dim)\n\n    def forward(self, x, mask):\n        mask_pad = (mask).long().bool().unsqueeze(1)\n        mask_pad = ~(mask_pad.permute(0, 2, 1) * mask_pad)\n        mask_flat = mask.view(-1).bool()\n        bs, slen, nfeats = x.shape\n\n        # MHSA\n        residual = x\n        x = x * self.scale_mhsa + self.bias_mhsa\n        x = residual + self.mhsa(x, mask_pad)\n\n        # Skip padding for layer norm\n        x_skip = x.reshape(-1, x.shape[-1])\n        x = x_skip[mask_flat].unsqueeze(0)\n        x = self.ln_mhsa(x)\n\n        # FFN after MHSA\n        residual = x\n        x = x * self.scale_ff_mhsa + self.bias_ff_mhsa\n        x = residual + self.ff_mhsa(x)\n        x = self.ln_ff_mhsa(x)\n\n        # Restore shape\n        x_skip[mask_flat] = x[0].to(dtype=x_skip.dtype)\n        x = x_skip.reshape(bs, slen, nfeats)\n\n        # Conv\n        residual = x\n        x = x * self.scale_conv + self.bias_conv\n        x = residual + self.conv(x, mask_pad=mask.bool().unsqueeze(1))\n\n        # Skip padding for layer norm\n        x_skip = x.reshape(-1, x.shape[-1])\n        x = x_skip[mask_flat].unsqueeze(0)\n        x = self.ln_conv(x)\n\n        # FFN after Conv\n        residual = x\n        x = x * self.scale_ff_conv + self.bias_ff_conv\n        x = residual + self.ff_conv(x)\n        x = self.ln_ff_conv(x)\n\n        # Restore shape\n        x_skip[mask_flat] = x[0].to(dtype=x_skip.dtype)\n        x = x_skip.reshape(bs, slen, nfeats)\n\n        return x\n\nclass SqueezeformerEncoder(nn.Module):\n    \"\"\"\n    Stack of Squeezeformer blocks\n    \"\"\"\n    def __init__(\n        self,\n        input_dim: int = 80,\n        encoder_dim: int = 512,\n        num_layers: int = 16,\n        num_attention_heads: int = 8,\n        feed_forward_expansion_factor: int = 4,\n        conv_expansion_factor: int = 2,\n        input_dropout_p: float = 0.0,\n        feed_forward_dropout_p: float = 0.0,\n        attention_dropout_p: float = 0.0,\n        conv_dropout_p: float = 0.0,\n        conv_kernel_size: int = 31,\n        use_bn: bool = True,\n    ):\n        super(SqueezeformerEncoder, self).__init__()\n        self.num_layers = num_layers\n\n        self.blocks = nn.ModuleList()\n        for idx in range(num_layers):\n            self.blocks.append(\n                SqueezeformerBlock(\n                    encoder_dim=encoder_dim,\n                    num_attention_heads=num_attention_heads,\n                    feed_forward_expansion_factor=feed_forward_expansion_factor,\n                    conv_expansion_factor=conv_expansion_factor,\n                    feed_forward_dropout_p=feed_forward_dropout_p,\n                    attention_dropout_p=attention_dropout_p,\n                    conv_dropout_p=conv_dropout_p,\n                    conv_kernel_size=conv_kernel_size,\n                    use_bn=use_bn,\n                )\n            )\n\n    def forward(self, x: Tensor, mask: Tensor):\n        for idx, block in enumerate(self.blocks):\n            x = block(x, mask)\n        return x\n\nfrom timm.layers.norm_act import BatchNormAct2d\n\nclass FeatureExtractor(nn.Module):\n    def __init__(self,\n                 n_landmarks,\n                 out_dim):\n        super().__init__()   \n\n        self.in_channels = in_channels = (32//2) * n_landmarks\n        self.stem_linear = nn.Linear(in_channels,out_dim,bias=False)\n        self.stem_bn = nn.BatchNorm1d(out_dim, momentum=0.95)\n        self.conv_stem = nn.Conv2d(4, 32, kernel_size=(3, 3), stride=(1, 2), padding=(1, 1), bias=False)\n        self.bn_conv = BatchNormAct2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True,act_layer = nn.SiLU,drop_layer=None)\n        \n    def forward(self, data, mask):\n        xc = data.permute(0,2,1,3)  # B,C,T,F\n        xc = self.conv_stem(xc)\n        xc = self.bn_conv(xc)\n        xc = xc.permute(0,2,3,1)\n        xc = xc.reshape(*data.shape[:2], -1)\n        \n        m = mask.to(torch.bool)  \n        x = self.stem_linear(xc)\n        \n        # Batchnorm without pads\n        bs,slen,nfeat = x.shape\n        x = x.view(-1, nfeat)\n        x_bn = x[mask.view(-1)==1].unsqueeze(0)\n        x_bn = self.stem_bn(x_bn.permute(0,2,1)).permute(0,2,1)\n        x[mask.view(-1)==1] = x_bn[0]\n        x = x.view(bs,slen,nfeat)\n        # Padding mask\n        x = x.masked_fill(~mask.bool().unsqueeze(-1), 0.0)\n        \n        return x\n\nclass MABeFeatureExtractor(nn.Module):\n    def __init__(self, input_dim=708, encoder_dim=144, dropout=0.0):\n        super().__init__()\n\n        self.input_dim = input_dim\n        self.encoder_dim = encoder_dim\n\n        # Project input features to encoder dimension\n        self.input_proj = nn.Linear(input_dim, encoder_dim)\n        self.input_norm = nn.LayerNorm(encoder_dim)\n        self.input_dropout = nn.Dropout(dropout)\n\n        # Optional: Add a small CNN for local temporal patterns\n        self.use_conv = True\n        if self.use_conv:\n            self.conv1 = nn.Conv1d(encoder_dim, encoder_dim, kernel_size=3, padding=1)\n            self.conv2 = nn.Conv1d(encoder_dim, encoder_dim, kernel_size=3, padding=1)\n            self.conv_norm = nn.LayerNorm(encoder_dim)\n            self.conv_dropout = nn.Dropout(dropout)\n\n    def forward(self, x, mask):\n        \"\"\"\n        Args:\n            x: Input features (batch, seq_len, 708)\n            mask: Attention mask (batch, seq_len)\n        Returns:\n            Encoded features (batch, seq_len, encoder_dim)\n        \"\"\"\n        # Project to encoder dimension\n        x = self.input_proj(x)\n        x = self.input_norm(x)\n        x = self.input_dropout(x)\n\n        if self.use_conv:\n            # Apply temporal convolutions\n            x_conv = x.transpose(1, 2)  # (batch, encoder_dim, seq_len)\n            x_conv = F.relu(self.conv1(x_conv))\n            x_conv = F.relu(self.conv2(x_conv))\n            x_conv = x_conv.transpose(1, 2)  # (batch, seq_len, encoder_dim)\n\n            # Residual connection\n            x = x + self.conv_dropout(self.conv_norm(x_conv))\n\n        # Apply mask\n        x = x.masked_fill(~mask.bool().unsqueeze(-1), 0.0)\n\n        return x\n\nimport torch.nn.functional as F\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch_geometric.nn import TransformerConv, LayerNorm\n\nclass SpatialMouseGNN(nn.Module):\n    def __init__(self, in_channels, out_channels, num_mice=4, heads=4, dropout=0.1):\n        super().__init__()\n        self.num_mice = num_mice\n        # 1. Project Raw Features (Crucial Step)\n        # We need to turn coordinates into \"semantic\" embeddings before Attention\n        self.embedding = nn.Linear(in_channels, out_channels)\n        \n        # 2. TransformerConv Layers\n        # We use beta=True for a gating mechanism (improves deep GNNs)\n        self.conv1 = TransformerConv(out_channels, out_channels // heads, heads=heads, \n                                     dropout=dropout, beta=True)\n        self.norm1 = LayerNorm(out_channels) # LayerNorm is often more stable for Transformers\n        \n        self.conv2 = TransformerConv(out_channels, out_channels // heads, heads=heads, \n                                     dropout=dropout, beta=True)\n        self.norm2 = LayerNorm(out_channels)\n        \n        self.dropout = nn.Dropout(dropout)\n\n    def _get_fully_connected_edge_index(self, batch_size_time, device):\n        \"\"\"\n        Creates edges so every mouse connects to every other mouse within the same frame.\n        \"\"\"\n        # Base edges for one frame of 4 mice (0-3)\n        # 0->1, 0->2, 0->3, 1->0, ...\n        base_src = torch.tensor([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], device=device)\n        base_dst = torch.tensor([1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2], device=device)\n        # If you want self-loops (mice look at themselves), add them here, \n        # but TransformerConv usually adds them automatically or handles them separately.\n\n        # Repeat for the whole batch\n        # We need to offset indices by num_mice (4) for each subsequent frame\n        offsets = torch.arange(batch_size_time, device=device) * self.num_mice\n        \n        # Broadcasting to create the massive edge list\n        # shape: [Edges_Per_Frame, Batch_Time]\n        src = base_src.unsqueeze(1) + offsets.unsqueeze(0)\n        dst = base_dst.unsqueeze(1) + offsets.unsqueeze(0)\n        \n        edge_index = torch.stack([src.flatten(), dst.flatten()], dim=0)\n        return edge_index\n\n    def forward(self, x):\n        \"\"\"\n        Input: [Batch, Frames, Mice, Features]\n        Output: [Batch, Frames, Mice, Out_Channels]\n        \"\"\"\n        B, T, M, D = x.shape\n        \n        # Flatten Batch and Time: treat every frame as an independent graph\n        # New shape: [Total_Graphs * Mice, Features]\n        x_flat = x.view(B * T * M, D)\n        \n        # 1. Linear Projection\n        x_emb = F.relu(self.embedding(x_flat))\n        \n        # 2. Create Edges on the fly\n        # (Optimized: we only calculate this once per forward pass structure)\n        edge_index = self._get_fully_connected_edge_index(B * T, x.device)\n        \n        # 3. GNN Layer 1\n        # TransformerConv expects [Num_Nodes, Dim]\n        h = self.conv1(x_emb, edge_index)\n        h = self.norm1(h)\n        h = F.relu(h)\n        h = self.dropout(h)\n        \n        # 4. GNN Layer 2\n        h = self.conv2(h, edge_index)\n        h = self.norm2(h)\n        h = F.relu(h)\n        \n        # 5. Reshape back (NO POOLING)\n        # We want to keep the mice separate for the Squeezeformer\n        return h.view(B, T, M, -1)\n\n\nclass BehaviorClassificationHead(nn.Module):\n\n    def __init__(self, encoder_dim=144, num_pairs=16, num_actions=39, dropout=0.0):\n        super().__init__()\n\n        self.num_pairs = num_pairs  # Number of mouse pairs (e.g., 4 mice -> 16 directed pairs)\n        self.num_actions = num_actions  # Number of behavior classes\n\n        self.shared_proj = nn.Sequential(\n            nn.Linear(encoder_dim, encoder_dim * 2),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(encoder_dim * 2, encoder_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout)\n        )\n        self.classifier = nn.Linear(encoder_dim, num_pairs * num_actions)\n\n    def forward(self, x):\n        batch_size, seq_len, _ = x.shape\n        x = self.shared_proj(x)  # (batch, seq_len, encoder_dim)\n        logits = self.classifier(x)  # (batch, seq_len, num_pairs * num_actions)\n        logits = logits.view(batch_size, seq_len, self.num_pairs, self.num_actions)\n\n        return logits\n\nclass PairwiseBehaviorHead(nn.Module):\n    def __init__(self, encoder_dim, num_classes=39, dropout=0.0):\n        super().__init__()\n        \n        self.in_features = encoder_dim * 2 \n        \n        self.layers = nn.Sequential(\n            nn.Linear(self.in_features, encoder_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            # Output is logits for ONE pair\n            nn.Linear(encoder_dim, num_classes)\n        )\n\n    def forward(self, x):\n        \"\"\"\n        Args:\n            x: [Batch, Time, Mice, Dim] - The unflattened Squeezeformer output\n        Returns:\n            logits: [Batch, Time, Num_Pairs (16), Num_Classes]\n        \"\"\"\n        B, T, M, D = x.shape\n        \n        # We need to form 16 pairs. \n        # Efficient way using broadcasting:\n        \n        # 1. Expand M1 to represent the \"Source\" mouse\n        # Shape: [B, T, M, 1, D] -> Repeat on dim 3 -> [B, T, M, M, D]\n        m1 = x.unsqueeze(3).expand(-1, -1, -1, M, -1)\n        \n        # 2. Expand M2 to represent the \"Target\" mouse\n        # Shape: [B, T, 1, M, D] -> Repeat on dim 2 -> [B, T, M, M, D]\n        m2 = x.unsqueeze(2).expand(-1, -1, M, -1, -1)\n        \n        # 3. Concatenate to get pair features\n        # Shape: [B, T, M, M, 2*D]\n        pair_features = torch.cat([m1, m2], dim=-1)\n        \n        # 4. Flatten the mouse dimensions to get the list of 16 pairs\n        # Shape: [B, T, M*M, 2*D] -> [B, T, 16, 2*D]\n        pair_features_flat = pair_features.view(B, T, M*M, -1)\n        \n        # 5. Classify\n        logits = self.layers(pair_features_flat)\n        \n        return logits\n\nclass NetGNN(nn.Module):\n    \"\"\"\n    Squeezeformer model for MABe mouse behavior detection\n    \"\"\"\n    def __init__(self, cfg):\n        super(NetGNN, self).__init__()\n        self.cfg = cfg\n        # Model dimensions\n        self.encoder_dim = cfg.encoder_config.encoder_dim\n        self.num_pairs = 16  # Number of mouse pairs\n        self.num_actions = 39  # Number of behavior classes (including no_action)\n\n        self.gnn = SpatialMouseGNN(\n            in_channels=cfg.per_mouse_feature_dim,\n            out_channels=cfg.encoder_config.encoder_dim,\n            num_mice=4,\n            heads=4,\n            dropout=cfg.encoder_config.input_dropout_p\n        )\n\n        # Feature extractor\n        if self.cfg.cnn_extractor:\n            self.feature_extractor = FeatureExtractor(\n                n_landmarks=cfg.per_mouse_feature_dim // 4,\n                out_dim=cfg.encoder_config.encoder_dim)\n        else:\n            self.feature_extractor = MABeFeatureExtractor(\n                input_dim=cfg.feature_dim,\n                encoder_dim=self.encoder_dim,\n                dropout=cfg.encoder_config.input_dropout_p\n            )\n\n        self.encoder = SqueezeformerEncoder(\n            input_dim=self.encoder_dim,\n            encoder_dim=self.encoder_dim,\n            num_layers=cfg.encoder_config.num_layers,\n            num_attention_heads=cfg.encoder_config.num_attention_heads,\n            feed_forward_expansion_factor=cfg.encoder_config.feed_forward_expansion_factor,\n            conv_expansion_factor=cfg.encoder_config.conv_expansion_factor,\n            input_dropout_p=cfg.encoder_config.input_dropout_p,\n            feed_forward_dropout_p=cfg.encoder_config.feed_forward_dropout_p,\n            attention_dropout_p=cfg.encoder_config.attention_dropout_p,\n            conv_dropout_p=cfg.encoder_config.conv_dropout_p,\n            conv_kernel_size=cfg.encoder_config.conv_kernel_size,\n            use_bn = cfg.use_bn\n        )\n\n        self.classifier = BehaviorClassificationHead(\n            encoder_dim=self.encoder_dim,\n            num_pairs=self.num_pairs,\n            num_actions=self.num_actions,\n            dropout=cfg.encoder_config.feed_forward_dropout_p\n        )\n\n        self.classifier = PairwiseBehaviorHead(\n            encoder_dim=self.encoder_dim,\n            num_classes=self.num_actions,\n            dropout=cfg.encoder_config.feed_forward_dropout_p\n        )\n\n        if hasattr(cfg, 'class_weights') and cfg.class_weights is not None:\n            class_weights = torch.tensor(cfg.class_weights)\n        else:\n            pos_weight = cfg.pos_weight if hasattr(cfg, 'pos_weight') else 50.0 \n            class_weights = torch.full((self.num_actions,), pos_weight) # High weight for all\n            no_action_idx = cfg.action_id_map.get('no_action', -1)\n            if no_action_idx != -1:\n                class_weights[no_action_idx] = 1.0 # Low weight for no_action\n            else:\n                # Fallback if cfg is wrong, assume last class is no_action\n                class_weights[-1] = 1.0\n\n        self.register_buffer('class_weights', class_weights)\n\n        # Use CrossEntropyLoss for multi-class classification\n        self.loss_fn = nn.CrossEntropyLoss(weight=self.class_weights, reduction='none')\n        self.unweighted_ce_fn = nn.CrossEntropyLoss(reduction='none')\n\n        # Optional: Multi-class focal loss\n        self.use_focal_loss = cfg.use_focal_loss if hasattr(cfg, 'use_focal_loss') else False\n        self.focal_gamma = cfg.focal_gamma if hasattr(cfg, 'focal_gamma') else 2.0\n\n        # Training settings\n        self.return_logits = cfg.return_logits if hasattr(cfg, 'return_logits') else False\n\n        print(f'Model initialized with {count_parameters(self):,} trainable parameters')\n        # print(f'Loss: {\"Multi-class Focal\" if self.use_focal_loss else \"Weighted CrossEntropy\"}')\n        # print(f'Class weights: {self.class_weights.cpu().numpy()}')\n\n    def forward(self, batch):\n        mask = batch['input_mask'].long()  # (batch, seq_len)\n        mask_for_encoder = mask.repeat_interleave(4, dim=0)\n\n        if self.cfg.cnn_extractor:\n            x = batch['input_mice']\n            x = self.gnn(x)  # (batch, seq_len, num_mice, encoder_dim)\n            B, T, M, C = x.shape\n            x = x.permute(0,2,1,3).reshape(B * M, T, C)  # (batch*num_mice, seq_len, encoder_dim)\n        else:\n            x = batch['input']  # (batch, seq_len, 708)\n            B, T, _ = x.shape\n            x = self.feature_extractor(x, mask)  # (batch, seq_len, encoder_dim)\n\n\n        # Encode\n        x = self.encoder(x, mask_for_encoder)  # (batch, seq_len, encoder_dim)\n        x = x.reshape(B, M, T, C).permute(0,2,1,3)  # (batch, seq_len, num_mice, encoder_dim)\n\n        # Classify\n        logits = self.classifier(x)  # (batch, seq_len, num_pairs, num_actions)\n\n        # APPLY BEHAVIOR MASK if provided (for evaluation)\n        if 'behavior_mask' in batch and batch['behavior_mask'] is not None:\n            behavior_mask = batch['behavior_mask']  # (batch, num_pairs, num_actions)\n            # Expand to match logits shape\n            behavior_mask_expanded = behavior_mask.unsqueeze(1).expand_as(logits)\n            # For multi-class, we need to ensure at least one action is valid per pair\n            # Set masked positions to very negative value\n            logits = torch.where(behavior_mask_expanded.bool(), logits,\n                                torch.tensor(-1e10, dtype=logits.dtype, device=logits.device) if logits.dtype == torch.float32\n                                else torch.tensor(-1e4, dtype=logits.dtype, device=logits.device))\n\n        output = {}\n\n        # Calculate loss if labels provided\n        if 'labels' in batch and batch['labels'] is not None:\n            labels = batch['labels']\n\n            # Convert one-hot to class indices if needed\n            if labels.dim() == 4:  # One-hot encoded (batch, seq_len, num_pairs, num_actions)\n                # Convert to class indices\n                labels = torch.argmax(labels, dim=-1)  # (batch, seq_len, num_pairs)\n\n            # Reshape for loss calculation\n            batch_size, seq_len, num_pairs, num_actions = logits.shape\n            logits_flat = logits.reshape(-1, num_actions)  # (batch*seq*pairs, num_actions)\n            labels_flat = labels.reshape(-1)  # (batch*seq*pairs,)\n\n            # Calculate loss\n            if self.use_focal_loss:\n                # Multi-class focal loss\n                unweighted_ce = self.unweighted_ce_fn(logits_flat, labels_flat)\n                pt = torch.exp(-unweighted_ce)\n                alpha = self.class_weights[labels_flat]  # Gather alpha for true classes\n                focal_loss = alpha * (1 - pt) ** self.focal_gamma * unweighted_ce\n                loss = focal_loss\n            else:\n                # Standard weighted cross-entropy\n                loss = self.loss_fn(logits_flat, labels_flat)\n\n            # Reshape loss back\n            loss = loss.view(batch_size, seq_len, num_pairs)\n\n            mask_expanded = mask.unsqueeze(-1).expand_as(loss)\n            loss = loss * mask_expanded\n\n            # Average over valid positions\n            valid_positions = mask_expanded.sum()\n            if valid_positions > 0:\n                loss = loss.sum() / valid_positions\n            else:\n                loss = loss.sum()\n            output['loss'] = loss\n\n        probs = torch.softmax(logits, dim=-1)  # Multi-class probabilities\n        output['predictions'] = probs\n        output['logits'] = logits\n\n        # Store additional info for evaluation\n        if 'video_id' in batch:\n            output['video_id'] = batch['video_id']\n        if 'start_frame' in batch:\n            output['start_frame'] = batch['start_frame']\n\n        return output","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T13:06:48.944854Z","iopub.execute_input":"2025-12-18T13:06:48.945373Z","iopub.status.idle":"2025-12-18T13:06:48.976600Z","shell.execute_reply.started":"2025-12-18T13:06:48.945352Z","shell.execute_reply":"2025-12-18T13:06:48.975894Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### MODEL INTERLEAVE","metadata":{}},{"cell_type":"code","source":"%%writefile model_inter.py\n\n\nimport torch\nfrom torch.nn import functional as F\nfrom torch import nn\nfrom typing import Tuple, Union, Optional\nfrom torch import Tensor\nimport math\nimport numpy as np\nfrom torch_geometric.nn import TransformerConv, LayerNorm\n\n\ndef count_parameters(model):\n    \"\"\"Count the number of trainable parameters in model\"\"\"\n    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n\nclass Swish(nn.Module):\n    \"\"\"Swish activation function\"\"\"\n    def __init__(self) -> None:\n        super(Swish, self).__init__()\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return inputs * inputs.sigmoid()\n\nclass GLU(nn.Module):\n    \"\"\"Gated Linear Unit activation\"\"\"\n    def __init__(self, dim: int) -> None:\n        super(GLU, self).__init__()\n        self.dim = dim\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        outputs, gate = inputs.chunk(2, dim=self.dim)\n        return outputs * gate.sigmoid()\n\nclass FeedForwardModule(nn.Module):\n    \"\"\"\n    Feed Forward Module with pre-norm residual units\n    \"\"\"\n    def __init__(\n        self,\n        encoder_dim: int = 512,\n        expansion_factor: int = 4,\n        dropout_p: float = 0.0,\n    ) -> None:\n        super(FeedForwardModule, self).__init__()\n\n        self.ffn1 = nn.Linear(encoder_dim, encoder_dim * expansion_factor, bias=True)\n        self.act = Swish()\n        self.do1 = nn.Dropout(p=dropout_p)\n        self.ffn2 = nn.Linear(encoder_dim * expansion_factor, encoder_dim, bias=True)\n        self.do2 = nn.Dropout(p=dropout_p)\n\n    def forward(self, x):\n        x = self.ffn1(x)\n        x = self.act(x)\n        x = self.do1(x)\n        x = self.ffn2(x)\n        x = self.do2(x)\n        return x\n\nclass RelPositionalEncoding(nn.Module):\n    \"\"\"\n    Relative positional encoding module for handling variable sequence lengths\n    \"\"\"\n    def __init__(self, d_model: int = 512, max_len: int = 5000) -> None:\n        super(RelPositionalEncoding, self).__init__()\n        self.d_model = d_model\n        self.pe = None\n        self.extend_pe(torch.tensor(0.0).expand(1, max_len))\n\n    def extend_pe(self, x):\n        if self.pe is not None:\n            if self.pe.size(1) >= x.size(1) * 2 - 1:\n                if self.pe.dtype != x.dtype or self.pe.device != x.device:\n                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)\n                return\n\n        pe_positive = torch.zeros(x.size(1), self.d_model)\n        pe_negative = torch.zeros(x.size(1), self.d_model)\n        position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)\n        div_term = torch.exp(\n            torch.arange(0, self.d_model, 2, dtype=torch.float32) * -(math.log(10000.0) / self.d_model)\n        )\n        pe_positive[:, 0::2] = torch.sin(position * div_term)\n        pe_positive[:, 1::2] = torch.cos(position * div_term)\n        pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)\n        pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)\n\n        pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)\n        pe_negative = pe_negative[1:].unsqueeze(0)\n        pe = torch.cat([pe_positive, pe_negative], dim=1)\n        self.pe = pe.to(device=x.device, dtype=x.dtype)\n\n    def forward(self, x: torch.Tensor):\n        self.extend_pe(x)\n        pos_emb = self.pe[\n            :,\n            self.pe.size(1) // 2 - x.size(1) + 1 : self.pe.size(1) // 2 + x.size(1),\n        ]\n        return pos_emb\n\nclass RelativeMultiHeadAttention(nn.Module):\n    \"\"\"\n    Multi-head attention with relative positional encoding\n    \"\"\"\n    def __init__(\n        self,\n        d_model: int = 512,\n        num_heads: int = 16,\n        dropout_p: float = 0.0,\n    ):\n        super(RelativeMultiHeadAttention, self).__init__()\n        assert d_model % num_heads == 0, \"d_model % num_heads should be zero.\"\n        self.d_model = d_model\n        self.d_head = int(d_model / num_heads)\n        self.num_heads = num_heads\n        self.sqrt_dim = math.sqrt(self.d_head)\n\n        self.query_proj = nn.Linear(d_model, d_model)\n        self.key_proj = nn.Linear(d_model, d_model)\n        self.value_proj = nn.Linear(d_model, d_model)\n        self.pos_proj = nn.Linear(d_model, d_model, bias=False)\n\n        self.dropout = nn.Dropout(p=dropout_p)\n        self.u_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))\n        self.v_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))\n        torch.nn.init.xavier_uniform_(self.u_bias)\n        torch.nn.init.xavier_uniform_(self.v_bias)\n\n        self.out_proj = nn.Linear(d_model, d_model)\n\n    def forward(\n        self,\n        query: Tensor,\n        key: Tensor,\n        value: Tensor,\n        pos_embedding: Tensor,\n        mask: Optional[Tensor] = None,\n    ) -> Tensor:\n        batch_size = value.size(0)\n\n        query = self.query_proj(query).view(batch_size, -1, self.num_heads, self.d_head)\n        key = self.key_proj(key).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)\n        value = self.value_proj(value).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)\n        pos_embedding = self.pos_proj(pos_embedding).view(batch_size, -1, self.num_heads, self.d_head)\n\n        content_score = torch.matmul((query + self.u_bias).transpose(1, 2), key.transpose(2, 3))\n        pos_score = torch.matmul((query + self.v_bias).transpose(1, 2), pos_embedding.permute(0, 2, 3, 1))\n        pos_score = self._relative_shift(pos_score)\n\n        score = (content_score + pos_score) / self.sqrt_dim\n\n        if mask is not None:\n            mask = mask.unsqueeze(1)\n            score.masked_fill_(mask, -1e9) if score.dtype == torch.float32 else score.masked_fill_(mask, -1e4)\n\n        attn = F.softmax(score, -1)\n        attn = self.dropout(attn)\n\n        context = torch.matmul(attn, value).transpose(1, 2)\n        context = context.contiguous().view(batch_size, -1, self.d_model)\n\n        return self.out_proj(context)\n\n    def _relative_shift(self, pos_score: Tensor) -> Tensor:\n        batch_size, num_heads, seq_length1, seq_length2 = pos_score.size()\n        zeros = pos_score.new_zeros(batch_size, num_heads, seq_length1, 1)\n        padded_pos_score = torch.cat([zeros, pos_score], dim=-1)\n\n        padded_pos_score = padded_pos_score.view(batch_size, num_heads, seq_length2 + 1, seq_length1)\n        pos_score = padded_pos_score[:, :, 1:].view_as(pos_score)[:, :, :, : seq_length2 // 2 + 1]\n\n        return pos_score\n\nclass MultiHeadedSelfAttentionModule(nn.Module):\n    \"\"\"\n    Self-attention module with relative positional encoding\n    \"\"\"\n    def __init__(self, d_model: int, num_heads: int, dropout_p: float = 0.0):\n        super(MultiHeadedSelfAttentionModule, self).__init__()\n        self.positional_encoding = RelPositionalEncoding(d_model)\n        self.attention = RelativeMultiHeadAttention(d_model, num_heads, dropout_p)\n        self.dropout = nn.Dropout(p=dropout_p)\n\n    def forward(self, inputs: Tensor, mask: Optional[Tensor] = None):\n        batch_size = inputs.size(0)\n        pos_embedding = self.positional_encoding(inputs)\n        pos_embedding = pos_embedding.repeat(batch_size, 1, 1)\n\n        outputs = self.attention(inputs, inputs, inputs, pos_embedding=pos_embedding, mask=mask)\n        return self.dropout(outputs)\n\nclass DepthwiseConv1d(nn.Module):\n    \"\"\"Depthwise 1D convolution\"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        kernel_size: int,\n        stride: int = 1,\n        padding: int = 0,\n        bias: bool = False,\n    ) -> None:\n        super(DepthwiseConv1d, self).__init__()\n        assert out_channels % in_channels == 0, \"out_channels should be constant multiple of in_channels\"\n        self.conv = nn.Conv1d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=kernel_size,\n            groups=in_channels,\n            stride=stride,\n            padding=padding,\n            bias=bias,\n        )\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return self.conv(inputs)\n\nclass PointwiseConv1d(nn.Module):\n    \"\"\"Pointwise 1D convolution (kernel size = 1)\"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        stride: int = 1,\n        padding: int = 0,\n        bias: bool = True,\n    ) -> None:\n        super(PointwiseConv1d, self).__init__()\n        self.conv = nn.Conv1d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=1,\n            stride=stride,\n            padding=padding,\n            bias=bias,\n        )\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return self.conv(inputs)\n\nclass ConvModule(nn.Module):\n    \"\"\"\n    Convolution module with pointwise conv -> GLU -> depthwise conv -> normalization -> activation\n    \"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        kernel_size: int = 31,\n        expansion_factor: int = 2,\n        dropout_p: float = 0.0,\n        use_bn: bool = True,\n    ) -> None:\n        super(ConvModule, self).__init__()\n        assert (kernel_size - 1) % 2 == 0, \"kernel_size should be a odd number for 'SAME' padding\"\n        assert expansion_factor == 2, \"Currently, Only Supports expansion_factor 2\"\n\n        self.pw_conv_1 = PointwiseConv1d(in_channels, in_channels * expansion_factor, stride=1, padding=0, bias=True)\n        self.act1 = GLU(dim=1)\n        self.dw_conv = DepthwiseConv1d(in_channels, in_channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2)\n        self.bn = nn.BatchNorm1d(in_channels)\n        self.inorm = nn.InstanceNorm1d(in_channels, affine=True)\n        self.act2 = Swish()\n        self.pw_conv_2 = PointwiseConv1d(in_channels, in_channels, stride=1, padding=0, bias=True)\n        self.do = nn.Dropout(p=dropout_p)\n        self.use_bn = use_bn\n\n    def forward(self, x, mask_pad):\n        # Transpose for conv operations [B, T, C]\n        x = x.transpose(1, 2)\n        if mask_pad.size(2) > 0:  # time > 0\n            x = x.masked_fill(~mask_pad, 0.0)\n\n        x = self.pw_conv_1(x)\n        x = self.act1(x)\n        x = self.dw_conv(x)\n\n        if self.use_bn:\n            # Apply batch norm only to non-padded positions\n            x_bn = x.permute(0,2,1).reshape(-1, x.shape[1])\n            mask_bn = mask_pad.view(-1)\n            x_bn[mask_bn] = self.bn(x_bn[mask_bn])\n            x = x_bn.view(x.permute(0,2,1).shape).permute(0,2,1)\n        else:    \n            x = self.inorm(x)\n\n        x = self.act2(x)\n        x = self.pw_conv_2(x)\n        x = self.do(x)\n\n        # Mask batch padding again\n        if mask_pad.size(2) > 0:  # time > 0\n            x = x.masked_fill(~mask_pad, 0.0)\n        x = x.transpose(1, 2)\n        return x\n\ndef make_scale(encoder_dim):\n    \"\"\"Create learnable scale and bias parameters\"\"\"\n    scale = torch.nn.Parameter(torch.tensor([1.] * encoder_dim)[None, None, :])\n    bias = torch.nn.Parameter(torch.tensor([0.] * encoder_dim)[None, None, :])\n    return scale, bias\n\nclass SqueezeformerBlock(nn.Module):\n    \"\"\"\n    Squeezeformer block: MHSA -> FFN -> Conv -> FFN with residual connections\n    \"\"\"\n    def __init__(\n        self,\n        encoder_dim: int = 512,\n        num_attention_heads: int = 8,\n        feed_forward_expansion_factor: int = 4,\n        conv_expansion_factor: int = 2,\n        feed_forward_dropout_p: float = 0.0,\n        attention_dropout_p: float = 0.1,\n        conv_dropout_p: float = 0.1,\n        conv_kernel_size: int = 31,\n        use_bn: bool = True,\n    ):\n        super(SqueezeformerBlock, self).__init__()\n\n        self.scale_mhsa, self.bias_mhsa = make_scale(encoder_dim)\n        self.scale_ff_mhsa, self.bias_ff_mhsa = make_scale(encoder_dim)\n        self.scale_conv, self.bias_conv = make_scale(encoder_dim)\n        self.scale_ff_conv, self.bias_ff_conv = make_scale(encoder_dim)\n\n        self.mhsa = MultiHeadedSelfAttentionModule(\n            d_model=encoder_dim,\n            num_heads=num_attention_heads,\n            dropout_p=attention_dropout_p,\n        )\n        self.ln_mhsa = nn.LayerNorm(encoder_dim)\n        self.ff_mhsa = FeedForwardModule(\n            encoder_dim=encoder_dim,\n            expansion_factor=feed_forward_expansion_factor,\n            dropout_p=feed_forward_dropout_p,\n        )\n        self.ln_ff_mhsa = nn.LayerNorm(encoder_dim)\n        self.conv = ConvModule(\n            in_channels=encoder_dim,\n            kernel_size=conv_kernel_size,\n            expansion_factor=conv_expansion_factor,\n            dropout_p=conv_dropout_p,\n            use_bn=use_bn,\n        )\n        self.ln_conv = nn.LayerNorm(encoder_dim)\n        self.ff_conv = FeedForwardModule(\n            encoder_dim=encoder_dim,\n            expansion_factor=feed_forward_expansion_factor,\n            dropout_p=feed_forward_dropout_p,\n        )\n        self.ln_ff_conv = nn.LayerNorm(encoder_dim)\n\n    def forward(self, x, mask):\n        mask_pad = (mask).long().bool().unsqueeze(1)\n        mask_pad = ~(mask_pad.permute(0, 2, 1) * mask_pad)\n        mask_flat = mask.view(-1).bool()\n        bs, slen, nfeats = x.shape\n\n        # MHSA\n        residual = x\n        x = x * self.scale_mhsa + self.bias_mhsa\n        x = residual + self.mhsa(x, mask_pad)\n\n        # Skip padding for layer norm\n        x_skip = x.view(-1, x.shape[-1])\n        x = x_skip[mask_flat].unsqueeze(0)\n        x = self.ln_mhsa(x)\n\n        # FFN after MHSA\n        residual = x\n        x = x * self.scale_ff_mhsa + self.bias_ff_mhsa\n        x = residual + self.ff_mhsa(x)\n        x = self.ln_ff_mhsa(x)\n\n        # Restore shape\n        x_skip[mask_flat] = x[0].to(dtype=x_skip.dtype)\n        x = x_skip.view(bs, slen, nfeats)\n\n        # Conv\n        residual = x\n        x = x * self.scale_conv + self.bias_conv\n        x = residual + self.conv(x, mask_pad=mask.bool().unsqueeze(1))\n\n        # Skip padding for layer norm\n        x_skip = x.view(-1, x.shape[-1])\n        x = x_skip[mask_flat].unsqueeze(0)\n        x = self.ln_conv(x)\n\n        # FFN after Conv\n        residual = x\n        x = x * self.scale_ff_conv + self.bias_ff_conv\n        x = residual + self.ff_conv(x)\n        x = self.ln_ff_conv(x)\n\n        # Restore shape\n        x_skip[mask_flat] = x[0].to(dtype=x_skip.dtype)\n        x = x_skip.view(bs, slen, nfeats)\n\n        return x\n\nclass SpatialMouseGNN(nn.Module):\n    def __init__(self, in_channels, out_channels, num_mice=4, heads=4, dropout=0.1):\n        super().__init__()\n        self.num_mice = num_mice\n        # 1. Project Raw Features (Crucial Step)\n        # Check if we need to project (if dim changes) or just use as is\n        self.proj_input = (in_channels != out_channels)\n        if self.proj_input:\n            self.embedding = nn.Linear(in_channels, out_channels)\n        \n        # 2. TransformerConv Layers\n        # We use beta=True for a gating mechanism (improves deep GNNs)\n        self.conv1 = TransformerConv(out_channels, out_channels // heads, heads=heads, \n                                     dropout=dropout, beta=True)\n        self.norm1 = LayerNorm(out_channels) # LayerNorm is often more stable for Transformers\n        \n        self.conv2 = TransformerConv(out_channels, out_channels // heads, heads=heads, \n                                     dropout=dropout, beta=True)\n        self.norm2 = LayerNorm(out_channels)\n        \n        self.dropout = nn.Dropout(dropout)\n\n    def _get_fully_connected_edge_index(self, batch_size_time, device):\n        \"\"\"\n        Creates edges so every mouse connects to every other mouse within the same frame.\n        \"\"\"\n        # Base edges for one frame of 4 mice (0-3)\n        # 0->1, 0->2, 0->3, 1->0, ...\n        base_src = torch.tensor([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], device=device)\n        base_dst = torch.tensor([1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2], device=device)\n        # If you want self-loops (mice look at themselves), add them here, \n        # but TransformerConv usually adds them automatically or handles them separately.\n\n        # Repeat for the whole batch\n        # We need to offset indices by num_mice (4) for each subsequent frame\n        offsets = torch.arange(batch_size_time, device=device) * self.num_mice\n        \n        # Broadcasting to create the massive edge list\n        # shape: [Edges_Per_Frame, Batch_Time]\n        src = base_src.unsqueeze(1) + offsets.unsqueeze(0)\n        dst = base_dst.unsqueeze(1) + offsets.unsqueeze(0)\n        \n        edge_index = torch.stack([src.flatten(), dst.flatten()], dim=0)\n        return edge_index\n\n    def forward(self, x):\n        \"\"\"\n        Input: [Batch, Frames, Mice, Features]\n        Output: [Batch, Frames, Mice, Out_Channels]\n        \"\"\"\n        B, T, M, D = x.shape\n        \n        # Flatten Batch and Time: treat every frame as an independent graph\n        # New shape: [Total_Graphs * Mice, Features]\n        x_flat = x.reshape(B * T * M, D)\n        \n        # 1. Linear Projection\n        if self.proj_input:\n            x_emb = F.relu(self.embedding(x_flat))\n        else:\n            x_emb = x_flat\n        \n        # 2. Create Edges on the fly\n        # (Optimized: we only calculate this once per forward pass structure)\n        edge_index = self._get_fully_connected_edge_index(B * T, x.device)\n        \n        # 3. GNN Layer 1\n        # TransformerConv expects [Num_Nodes, Dim]\n        h = self.conv1(x_emb, edge_index)\n        h = self.norm1(h)\n        h = F.relu(h)\n        h = self.dropout(h)\n        \n        # 4. GNN Layer 2\n        h = self.conv2(h, edge_index)\n        h = self.norm2(h)\n        h = F.relu(h)\n        \n        # 5. Reshape back (NO POOLING)\n        # We want to keep the mice separate for the Squeezeformer\n        return h.reshape(B, T, M, -1)\n\nclass SqueezeformerEncoder(nn.Module):\n    \"\"\"\n    Stack of Squeezeformer blocks with Interleaved GNN Layers\n    \"\"\"\n    def __init__(\n        self,\n        input_dim: int = 80,\n        encoder_dim: int = 512,\n        num_layers: int = 16,\n        num_attention_heads: int = 8,\n        feed_forward_expansion_factor: int = 4,\n        conv_expansion_factor: int = 2,\n        input_dropout_p: float = 0.0,\n        feed_forward_dropout_p: float = 0.0,\n        attention_dropout_p: float = 0.0,\n        conv_dropout_p: float = 0.0,\n        conv_kernel_size: int = 31,\n        use_bn: bool = True,\n        gnn_interval: int = 4, # Inject GNN every N layers\n        num_mice: int = 4,\n    ):\n        super(SqueezeformerEncoder, self).__init__()\n        self.num_layers = num_layers\n        self.gnn_interval = gnn_interval\n        self.num_mice = num_mice\n        self.encoder_dim = encoder_dim\n\n        self.blocks = nn.ModuleList()\n        self.gnn_layers = nn.ModuleList()\n\n        for idx in range(num_layers):\n            self.blocks.append(\n                SqueezeformerBlock(\n                    encoder_dim=encoder_dim,\n                    num_attention_heads=num_attention_heads,\n                    feed_forward_expansion_factor=feed_forward_expansion_factor,\n                    conv_expansion_factor=conv_expansion_factor,\n                    feed_forward_dropout_p=feed_forward_dropout_p,\n                    attention_dropout_p=attention_dropout_p,\n                    conv_dropout_p=conv_dropout_p,\n                    conv_kernel_size=conv_kernel_size,\n                    use_bn=use_bn,\n                )\n            )\n            \n            # Interleaved GNN Injection\n            if (idx + 1) % gnn_interval == 0:\n                self.gnn_layers.append(\n                    SpatialMouseGNN(\n                        in_channels=encoder_dim, \n                        out_channels=encoder_dim, \n                        num_mice=num_mice,\n                        heads=4, \n                        dropout=0.1\n                    )\n                )\n            else:\n                self.gnn_layers.append(None)\n\n    def forward(self, x: Tensor, mask: Tensor):\n        # x input: [Batch * Mice, Time, Dim]\n        # mask input: [Batch * Mice, Time]\n        BM, T, D = x.shape\n        M = self.num_mice\n        B = BM // M\n\n        for idx, (block, gnn) in enumerate(zip(self.blocks, self.gnn_layers)):\n            x = block(x, mask)\n\n            if gnn is not None:\n                # 1. Reshape to [Batch, Time, Mice, Dim]\n                x_spatial = x.view(B, M, T, D).permute(0, 2, 1, 3)\n                \n                # 2. Apply GNN with Residual Connection\n                gnn_out = gnn(x_spatial)\n                x_spatial = x_spatial + gnn_out\n                \n                # 3. Flatten back to [Batch*Mice, Time, Dim]\n                x = x_spatial.permute(0, 2, 1, 3).reshape(BM, T, D)\n                \n        return x\n\nfrom timm.layers.norm_act import BatchNormAct2d\n\nclass FeatureExtractor(nn.Module):\n    def __init__(self,\n                 n_landmarks,\n                 out_dim):\n        super().__init__()   \n\n        self.in_channels = in_channels = (32//2) * n_landmarks\n        self.stem_linear = nn.Linear(in_channels,out_dim,bias=False)\n        self.stem_bn = nn.BatchNorm1d(out_dim, momentum=0.95)\n        self.conv_stem = nn.Conv2d(4, 32, kernel_size=(3, 3), stride=(1, 2), padding=(1, 1), bias=False)\n        self.bn_conv = BatchNormAct2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True,act_layer = nn.SiLU,drop_layer=None)\n        \n    def forward(self, data, mask):\n        xc = data.permute(0,2,1,3)  # B,C,T,F\n        xc = self.conv_stem(xc)\n        xc = self.bn_conv(xc)\n        xc = xc.permute(0,2,3,1)\n        xc = xc.reshape(*data.shape[:2], -1)\n        \n        m = mask.to(torch.bool)  \n        x = self.stem_linear(xc)\n        \n        # Batchnorm without pads\n        bs,slen,nfeat = x.shape\n        x = x.view(-1, nfeat)\n        x_bn = x[mask.view(-1)==1].unsqueeze(0)\n        x_bn = self.stem_bn(x_bn.permute(0,2,1)).permute(0,2,1)\n        x[mask.view(-1)==1] = x_bn[0]\n        x = x.view(bs,slen,nfeat)\n        # Padding mask\n        x = x.masked_fill(~mask.bool().unsqueeze(-1), 0.0)\n        \n        return x\n\nclass MABeFeatureExtractor(nn.Module):\n    def __init__(self, input_dim=708, encoder_dim=144, dropout=0.0):\n        super().__init__()\n\n        self.input_dim = input_dim\n        self.encoder_dim = encoder_dim\n\n        # Project input features to encoder dimension\n        self.input_proj = nn.Linear(input_dim, encoder_dim)\n        self.input_norm = nn.LayerNorm(encoder_dim)\n        self.input_dropout = nn.Dropout(dropout)\n\n        # Optional: Add a small CNN for local temporal patterns\n        self.use_conv = True\n        if self.use_conv:\n            self.conv1 = nn.Conv1d(encoder_dim, encoder_dim, kernel_size=3, padding=1)\n            self.conv2 = nn.Conv1d(encoder_dim, encoder_dim, kernel_size=3, padding=1)\n            self.conv_norm = nn.LayerNorm(encoder_dim)\n            self.conv_dropout = nn.Dropout(dropout)\n\n    def forward(self, x, mask):\n        \"\"\"\n        Args:\n            x: Input features (batch, seq_len, 708)\n            mask: Attention mask (batch, seq_len)\n        Returns:\n            Encoded features (batch, seq_len, encoder_dim)\n        \"\"\"\n        # Project to encoder dimension\n        x = self.input_proj(x)\n        x = self.input_norm(x)\n        x = self.input_dropout(x)\n\n        if self.use_conv:\n            # Apply temporal convolutions\n            x_conv = x.transpose(1, 2)  # (batch, encoder_dim, seq_len)\n            x_conv = F.relu(self.conv1(x_conv))\n            x_conv = F.relu(self.conv2(x_conv))\n            x_conv = x_conv.transpose(1, 2)  # (batch, seq_len, encoder_dim)\n\n            # Residual connection\n            x = x + self.conv_dropout(self.conv_norm(x_conv))\n\n        # Apply mask\n        x = x.masked_fill(~mask.bool().unsqueeze(-1), 0.0)\n\n        return x\n\nclass BehaviorClassificationHead(nn.Module):\n\n    def __init__(self, encoder_dim=144, num_pairs=16, num_actions=39, dropout=0.0):\n        super().__init__()\n\n        self.num_pairs = num_pairs  # Number of mouse pairs (e.g., 4 mice -> 16 directed pairs)\n        self.num_actions = num_actions  # Number of behavior classes\n\n        self.shared_proj = nn.Sequential(\n            nn.Linear(encoder_dim, encoder_dim * 2),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(encoder_dim * 2, encoder_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout)\n        )\n        self.classifier = nn.Linear(encoder_dim, num_pairs * num_actions)\n\n    def forward(self, x):\n        batch_size, seq_len, _ = x.shape\n        x = self.shared_proj(x)  # (batch, seq_len, encoder_dim)\n        logits = self.classifier(x)  # (batch, seq_len, num_pairs * num_actions)\n        logits = logits.view(batch_size, seq_len, self.num_pairs, self.num_actions)\n\n        return logits\n\nclass PairwiseBehaviorHead(nn.Module):\n    def __init__(self, encoder_dim, num_classes=39, dropout=0.0):\n        super().__init__()\n        \n        self.in_features = encoder_dim * 2 \n        \n        self.layers = nn.Sequential(\n            nn.Linear(self.in_features, encoder_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            # Output is logits for ONE pair\n            nn.Linear(encoder_dim, num_classes)\n        )\n\n    def forward(self, x):\n        \"\"\"\n        Args:\n            x: [Batch, Time, Mice, Dim] - The unflattened Squeezeformer output\n        Returns:\n            logits: [Batch, Time, Num_Pairs (16), Num_Classes]\n        \"\"\"\n        B, T, M, D = x.shape\n        \n        # We need to form 16 pairs. \n        # Efficient way using broadcasting:\n        \n        # 1. Expand M1 to represent the \"Source\" mouse\n        # Shape: [B, T, M, 1, D] -> Repeat on dim 3 -> [B, T, M, M, D]\n        m1 = x.unsqueeze(3).expand(-1, -1, -1, M, -1)\n        \n        # 2. Expand M2 to represent the \"Target\" mouse\n        # Shape: [B, T, 1, M, D] -> Repeat on dim 2 -> [B, T, M, M, D]\n        m2 = x.unsqueeze(2).expand(-1, -1, M, -1, -1)\n        \n        # 3. Concatenate to get pair features\n        # Shape: [B, T, M, M, 2*D]\n        pair_features = torch.cat([m1, m2], dim=-1)\n        \n        # 4. Flatten the mouse dimensions to get the list of 16 pairs\n        # Shape: [B, T, M*M, 2*D] -> [B, T, 16, 2*D]\n        pair_features_flat = pair_features.view(B, T, M*M, -1)\n        \n        # 5. Classify\n        logits = self.layers(pair_features_flat)\n        \n        return logits\n\nclass NetInter(nn.Module):\n    \"\"\"\n    Squeezeformer model for MABe mouse behavior detection\n    \"\"\"\n    def __init__(self, cfg):\n        super(NetInter, self).__init__()\n        self.cfg = cfg\n        # Model dimensions\n        self.encoder_dim = cfg.encoder_config.encoder_dim\n        self.num_pairs = 16  # Number of mouse pairs\n        self.num_actions = 39  # Number of behavior classes (including no_action)\n\n        self.gnn = SpatialMouseGNN(\n            in_channels=cfg.per_mouse_feature_dim,\n            out_channels=cfg.encoder_config.encoder_dim,\n            num_mice=4,\n            heads=4,\n            dropout=cfg.encoder_config.input_dropout_p\n        )\n\n        # Feature extractor\n        if self.cfg.cnn_extractor:\n            self.feature_extractor = FeatureExtractor(\n                n_landmarks=cfg.per_mouse_feature_dim // 4,\n                out_dim=cfg.encoder_config.encoder_dim)\n        else:\n            self.feature_extractor = MABeFeatureExtractor(\n                input_dim=cfg.feature_dim,\n                encoder_dim=self.encoder_dim,\n                dropout=cfg.encoder_config.input_dropout_p\n            )\n\n        self.encoder = SqueezeformerEncoder(\n            input_dim=self.encoder_dim,\n            encoder_dim=self.encoder_dim,\n            num_layers=cfg.encoder_config.num_layers,\n            num_attention_heads=cfg.encoder_config.num_attention_heads,\n            feed_forward_expansion_factor=cfg.encoder_config.feed_forward_expansion_factor,\n            conv_expansion_factor=cfg.encoder_config.conv_expansion_factor,\n            input_dropout_p=cfg.encoder_config.input_dropout_p,\n            feed_forward_dropout_p=cfg.encoder_config.feed_forward_dropout_p,\n            attention_dropout_p=cfg.encoder_config.attention_dropout_p,\n            conv_dropout_p=cfg.encoder_config.conv_dropout_p,\n            conv_kernel_size=cfg.encoder_config.conv_kernel_size,\n            use_bn = cfg.use_bn,\n            gnn_interval = 2, # Interleave GNN every 2 layers\n            num_mice = 4\n        )\n\n        self.classifier = BehaviorClassificationHead(\n            encoder_dim=self.encoder_dim,\n            num_pairs=self.num_pairs,\n            num_actions=self.num_actions,\n            dropout=cfg.encoder_config.feed_forward_dropout_p\n        )\n\n        self.classifier = PairwiseBehaviorHead(\n            encoder_dim=self.encoder_dim,\n            num_classes=self.num_actions,\n            dropout=cfg.encoder_config.feed_forward_dropout_p\n        )\n\n        if hasattr(cfg, 'class_weights') and cfg.class_weights is not None:\n            class_weights = torch.tensor(cfg.class_weights)\n        else:\n            pos_weight = cfg.pos_weight if hasattr(cfg, 'pos_weight') else 50.0 \n            class_weights = torch.full((self.num_actions,), pos_weight) # High weight for all\n            no_action_idx = cfg.action_id_map.get('no_action', -1)\n            if no_action_idx != -1:\n                class_weights[no_action_idx] = 1.0 # Low weight for no_action\n            else:\n                # Fallback if cfg is wrong, assume last class is no_action\n                class_weights[-1] = 1.0\n\n        self.register_buffer('class_weights', class_weights)\n\n        # Use CrossEntropyLoss for multi-class classification\n        self.loss_fn = nn.CrossEntropyLoss(weight=self.class_weights, reduction='none')\n        self.unweighted_ce_fn = nn.CrossEntropyLoss(reduction='none')\n\n        # Optional: Multi-class focal loss\n        self.use_focal_loss = cfg.use_focal_loss if hasattr(cfg, 'use_focal_loss') else False\n        self.focal_gamma = cfg.focal_gamma if hasattr(cfg, 'focal_gamma') else 2.0\n\n        # Training settings\n        self.return_logits = cfg.return_logits if hasattr(cfg, 'return_logits') else False\n\n        print(f'Model initialized with {count_parameters(self):,} trainable parameters')\n        # print(f'Loss: {\"Multi-class Focal\" if self.use_focal_loss else \"Weighted CrossEntropy\"}')\n        # print(f'Class weights: {self.class_weights.cpu().numpy()}')\n\n    def forward(self, batch):\n        mask = batch['input_mask'].long()  # (batch, seq_len)\n        mask_for_encoder = mask.repeat_interleave(4, dim=0)\n\n        if self.cfg.cnn_extractor:\n            x = batch['input_mice']\n            x = self.gnn(x)  # (batch, seq_len, num_mice, encoder_dim)\n            B, T, M, C = x.shape\n            x = x.permute(0,2,1,3).reshape(B * M, T, C)  # (batch*num_mice, seq_len, encoder_dim)\n        else:\n            x = batch['input']  # (batch, seq_len, 708)\n            B, T, _ = x.shape\n            x = self.feature_extractor(x, mask)  # (batch, seq_len, encoder_dim)\n\n\n        # Encode\n        x = self.encoder(x, mask_for_encoder)  # (batch, seq_len, encoder_dim)\n        x = x.reshape(B, M, T, C).permute(0,2,1,3)  # (batch, seq_len, num_mice, encoder_dim)\n\n        # Classify\n        logits = self.classifier(x)  # (batch, seq_len, num_pairs, num_actions)\n\n        # APPLY BEHAVIOR MASK if provided (for evaluation)\n        if 'behavior_mask' in batch and batch['behavior_mask'] is not None:\n            behavior_mask = batch['behavior_mask']  # (batch, num_pairs, num_actions)\n            # Expand to match logits shape\n            behavior_mask_expanded = behavior_mask.unsqueeze(1).expand_as(logits)\n            # For multi-class, we need to ensure at least one action is valid per pair\n            # Set masked positions to very negative value\n            logits = torch.where(behavior_mask_expanded.bool(), logits,\n                                torch.tensor(-1e10, dtype=logits.dtype, device=logits.device) if logits.dtype == torch.float32\n                                else torch.tensor(-1e4, dtype=logits.dtype, device=logits.device))\n\n        output = {}\n\n        # Calculate loss if labels provided\n        if 'labels' in batch and batch['labels'] is not None:\n            labels = batch['labels']\n\n            # Convert one-hot to class indices if needed\n            if labels.dim() == 4:  # One-hot encoded (batch, seq_len, num_pairs, num_actions)\n                # Convert to class indices\n                labels = torch.argmax(labels, dim=-1)  # (batch, seq_len, num_pairs)\n\n            # Reshape for loss calculation\n            batch_size, seq_len, num_pairs, num_actions = logits.shape\n            logits_flat = logits.reshape(-1, num_actions)  # (batch*seq*pairs, num_actions)\n            labels_flat = labels.reshape(-1)  # (batch*seq*pairs,)\n\n            # Calculate loss\n            if self.use_focal_loss:\n                # Multi-class focal loss\n                unweighted_ce = self.unweighted_ce_fn(logits_flat, labels_flat)\n                pt = torch.exp(-unweighted_ce)\n                alpha = self.class_weights[labels_flat]  # Gather alpha for true classes\n                focal_loss = alpha * (1 - pt) ** self.focal_gamma * unweighted_ce\n                loss = focal_loss\n            else:\n                # Standard weighted cross-entropy\n                loss = self.loss_fn(logits_flat, labels_flat)\n\n            # Reshape loss back\n            loss = loss.view(batch_size, seq_len, num_pairs)\n\n            mask_expanded = mask.unsqueeze(-1).expand_as(loss)\n            loss = loss * mask_expanded\n\n            # Average over valid positions\n            valid_positions = mask_expanded.sum()\n            if valid_positions > 0:\n                loss = loss.sum() / valid_positions\n            else:\n                loss = loss.sum()\n            output['loss'] = loss\n\n        probs = torch.softmax(logits, dim=-1)  # Multi-class probabilities\n        output['predictions'] = probs\n        output['logits'] = logits\n\n        # Store additional info for evaluation\n        if 'video_id' in batch:\n            output['video_id'] = batch['video_id']\n        if 'start_frame' in batch:\n            output['start_frame'] = batch['start_frame']\n\n        return output","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T13:06:48.977622Z","iopub.execute_input":"2025-12-18T13:06:48.977940Z","iopub.status.idle":"2025-12-18T13:06:49.007389Z","shell.execute_reply.started":"2025-12-18T13:06:48.977922Z","shell.execute_reply":"2025-12-18T13:06:49.006651Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### MODEL 238","metadata":{}},{"cell_type":"code","source":"%%writefile model_238.py\n\n\nimport torch\nfrom torch.nn import functional as F\nfrom torch import nn\nfrom typing import Tuple, Union, Optional\nfrom torch import Tensor\nimport math\nimport numpy as np\n\n\ndef count_parameters(model):\n    \"\"\"Count the number of trainable parameters in model\"\"\"\n    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n\nclass Swish(nn.Module):\n    \"\"\"Swish activation function\"\"\"\n    def __init__(self) -> None:\n        super(Swish, self).__init__()\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return inputs * inputs.sigmoid()\n\nclass GLU(nn.Module):\n    \"\"\"Gated Linear Unit activation\"\"\"\n    def __init__(self, dim: int) -> None:\n        super(GLU, self).__init__()\n        self.dim = dim\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        outputs, gate = inputs.chunk(2, dim=self.dim)\n        return outputs * gate.sigmoid()\n\nclass FeedForwardModule(nn.Module):\n    \"\"\"\n    Feed Forward Module with pre-norm residual units\n    \"\"\"\n    def __init__(\n        self,\n        encoder_dim: int = 512,\n        expansion_factor: int = 4,\n        dropout_p: float = 0.0,\n    ) -> None:\n        super(FeedForwardModule, self).__init__()\n\n        self.ffn1 = nn.Linear(encoder_dim, encoder_dim * expansion_factor, bias=True)\n        self.act = Swish()\n        self.do1 = nn.Dropout(p=dropout_p)\n        self.ffn2 = nn.Linear(encoder_dim * expansion_factor, encoder_dim, bias=True)\n        self.do2 = nn.Dropout(p=dropout_p)\n\n    def forward(self, x):\n        x = self.ffn1(x)\n        x = self.act(x)\n        x = self.do1(x)\n        x = self.ffn2(x)\n        x = self.do2(x)\n        return x\n\nclass RelPositionalEncoding(nn.Module):\n    \"\"\"\n    Relative positional encoding module for handling variable sequence lengths\n    \"\"\"\n    def __init__(self, d_model: int = 512, max_len: int = 5000) -> None:\n        super(RelPositionalEncoding, self).__init__()\n        self.d_model = d_model\n        self.pe = None\n        self.extend_pe(torch.tensor(0.0).expand(1, max_len))\n\n    def extend_pe(self, x):\n        if self.pe is not None:\n            if self.pe.size(1) >= x.size(1) * 2 - 1:\n                if self.pe.dtype != x.dtype or self.pe.device != x.device:\n                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)\n                return\n\n        pe_positive = torch.zeros(x.size(1), self.d_model)\n        pe_negative = torch.zeros(x.size(1), self.d_model)\n        position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)\n        div_term = torch.exp(\n            torch.arange(0, self.d_model, 2, dtype=torch.float32) * -(math.log(10000.0) / self.d_model)\n        )\n        pe_positive[:, 0::2] = torch.sin(position * div_term)\n        pe_positive[:, 1::2] = torch.cos(position * div_term)\n        pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)\n        pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)\n\n        pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)\n        pe_negative = pe_negative[1:].unsqueeze(0)\n        pe = torch.cat([pe_positive, pe_negative], dim=1)\n        self.pe = pe.to(device=x.device, dtype=x.dtype)\n\n    def forward(self, x: torch.Tensor):\n        self.extend_pe(x)\n        pos_emb = self.pe[\n            :,\n            self.pe.size(1) // 2 - x.size(1) + 1 : self.pe.size(1) // 2 + x.size(1),\n        ]\n        return pos_emb\n\nclass RelativeMultiHeadAttention(nn.Module):\n    \"\"\"\n    Multi-head attention with relative positional encoding\n    \"\"\"\n    def __init__(\n        self,\n        d_model: int = 512,\n        num_heads: int = 16,\n        dropout_p: float = 0.0,\n    ):\n        super(RelativeMultiHeadAttention, self).__init__()\n        assert d_model % num_heads == 0, \"d_model % num_heads should be zero.\"\n        self.d_model = d_model\n        self.d_head = int(d_model / num_heads)\n        self.num_heads = num_heads\n        self.sqrt_dim = math.sqrt(self.d_head)\n\n        self.query_proj = nn.Linear(d_model, d_model)\n        self.key_proj = nn.Linear(d_model, d_model)\n        self.value_proj = nn.Linear(d_model, d_model)\n        self.pos_proj = nn.Linear(d_model, d_model, bias=False)\n\n        self.dropout = nn.Dropout(p=dropout_p)\n        self.u_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))\n        self.v_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))\n        torch.nn.init.xavier_uniform_(self.u_bias)\n        torch.nn.init.xavier_uniform_(self.v_bias)\n\n        self.out_proj = nn.Linear(d_model, d_model)\n\n    def forward(\n        self,\n        query: Tensor,\n        key: Tensor,\n        value: Tensor,\n        pos_embedding: Tensor,\n        mask: Optional[Tensor] = None,\n    ) -> Tensor:\n        batch_size = value.size(0)\n\n        query = self.query_proj(query).view(batch_size, -1, self.num_heads, self.d_head)\n        key = self.key_proj(key).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)\n        value = self.value_proj(value).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)\n        pos_embedding = self.pos_proj(pos_embedding).view(batch_size, -1, self.num_heads, self.d_head)\n\n        content_score = torch.matmul((query + self.u_bias).transpose(1, 2), key.transpose(2, 3))\n        pos_score = torch.matmul((query + self.v_bias).transpose(1, 2), pos_embedding.permute(0, 2, 3, 1))\n        pos_score = self._relative_shift(pos_score)\n\n        score = (content_score + pos_score) / self.sqrt_dim\n\n        if mask is not None:\n            mask = mask.unsqueeze(1)\n            score.masked_fill_(mask, -1e9) if score.dtype == torch.float32 else score.masked_fill_(mask, -1e4)\n\n        attn = F.softmax(score, -1)\n        attn = self.dropout(attn)\n\n        context = torch.matmul(attn, value).transpose(1, 2)\n        context = context.contiguous().view(batch_size, -1, self.d_model)\n\n        return self.out_proj(context)\n\n    def _relative_shift(self, pos_score: Tensor) -> Tensor:\n        batch_size, num_heads, seq_length1, seq_length2 = pos_score.size()\n        zeros = pos_score.new_zeros(batch_size, num_heads, seq_length1, 1)\n        padded_pos_score = torch.cat([zeros, pos_score], dim=-1)\n\n        padded_pos_score = padded_pos_score.view(batch_size, num_heads, seq_length2 + 1, seq_length1)\n        pos_score = padded_pos_score[:, :, 1:].view_as(pos_score)[:, :, :, : seq_length2 // 2 + 1]\n\n        return pos_score\n\nclass MultiHeadedSelfAttentionModule(nn.Module):\n    \"\"\"\n    Self-attention module with relative positional encoding\n    \"\"\"\n    def __init__(self, d_model: int, num_heads: int, dropout_p: float = 0.0):\n        super(MultiHeadedSelfAttentionModule, self).__init__()\n        self.positional_encoding = RelPositionalEncoding(d_model)\n        self.attention = RelativeMultiHeadAttention(d_model, num_heads, dropout_p)\n        self.dropout = nn.Dropout(p=dropout_p)\n\n    def forward(self, inputs: Tensor, mask: Optional[Tensor] = None):\n        batch_size = inputs.size(0)\n        pos_embedding = self.positional_encoding(inputs)\n        pos_embedding = pos_embedding.repeat(batch_size, 1, 1)\n\n        outputs = self.attention(inputs, inputs, inputs, pos_embedding=pos_embedding, mask=mask)\n        return self.dropout(outputs)\n\nclass DepthwiseConv1d(nn.Module):\n    \"\"\"Depthwise 1D convolution\"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        kernel_size: int,\n        stride: int = 1,\n        padding: int = 0,\n        bias: bool = False,\n    ) -> None:\n        super(DepthwiseConv1d, self).__init__()\n        assert out_channels % in_channels == 0, \"out_channels should be constant multiple of in_channels\"\n        self.conv = nn.Conv1d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=kernel_size,\n            groups=in_channels,\n            stride=stride,\n            padding=padding,\n            bias=bias,\n        )\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return self.conv(inputs)\n\nclass PointwiseConv1d(nn.Module):\n    \"\"\"Pointwise 1D convolution (kernel size = 1)\"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        stride: int = 1,\n        padding: int = 0,\n        bias: bool = True,\n    ) -> None:\n        super(PointwiseConv1d, self).__init__()\n        self.conv = nn.Conv1d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=1,\n            stride=stride,\n            padding=padding,\n            bias=bias,\n        )\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return self.conv(inputs)\n\nclass ConvModule(nn.Module):\n    \"\"\"\n    Convolution module with pointwise conv -> GLU -> depthwise conv -> normalization -> activation\n    \"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        kernel_size: int = 31,\n        expansion_factor: int = 2,\n        dropout_p: float = 0.0,\n        use_bn: bool = True,\n    ) -> None:\n        super(ConvModule, self).__init__()\n        assert (kernel_size - 1) % 2 == 0, \"kernel_size should be a odd number for 'SAME' padding\"\n        assert expansion_factor == 2, \"Currently, Only Supports expansion_factor 2\"\n\n        self.pw_conv_1 = PointwiseConv1d(in_channels, in_channels * expansion_factor, stride=1, padding=0, bias=True)\n        self.act1 = GLU(dim=1)\n        self.dw_conv = DepthwiseConv1d(in_channels, in_channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2)\n        self.bn = nn.BatchNorm1d(in_channels)\n        self.inorm = nn.InstanceNorm1d(in_channels, affine=True)\n        self.act2 = Swish()\n        self.pw_conv_2 = PointwiseConv1d(in_channels, in_channels, stride=1, padding=0, bias=True)\n        self.do = nn.Dropout(p=dropout_p)\n        self.use_bn = use_bn\n\n    def forward(self, x, mask_pad):\n        # Transpose for conv operations [B, T, C]\n        x = x.transpose(1, 2)\n        if mask_pad.size(2) > 0:  # time > 0\n            x = x.masked_fill(~mask_pad, 0.0)\n\n        x = self.pw_conv_1(x)\n        x = self.act1(x)\n        x = self.dw_conv(x)\n\n        if self.use_bn:\n            # Apply batch norm only to non-padded positions\n            x_bn = x.permute(0,2,1).reshape(-1, x.shape[1])\n            mask_bn = mask_pad.reshape(-1)\n            x_bn[mask_bn] = self.bn(x_bn[mask_bn])\n            x = x_bn.reshape(x.permute(0,2,1).shape).permute(0,2,1)\n        else:    \n            x = self.inorm(x)\n\n        x = self.act2(x)\n        x = self.pw_conv_2(x)\n        x = self.do(x)\n\n        # Mask batch padding again\n        if mask_pad.size(2) > 0:  # time > 0\n            x = x.masked_fill(~mask_pad, 0.0)\n        x = x.transpose(1, 2)\n        return x\n\ndef make_scale(encoder_dim):\n    \"\"\"Create learnable scale and bias parameters\"\"\"\n    scale = torch.nn.Parameter(torch.tensor([1.] * encoder_dim)[None, None, :])\n    bias = torch.nn.Parameter(torch.tensor([0.] * encoder_dim)[None, None, :])\n    return scale, bias\n\nclass SqueezeformerBlock(nn.Module):\n    \"\"\"\n    Squeezeformer block: MHSA -> FFN -> Conv -> FFN with residual connections\n    \"\"\"\n    def __init__(\n        self,\n        encoder_dim: int = 512,\n        num_attention_heads: int = 8,\n        feed_forward_expansion_factor: int = 4,\n        conv_expansion_factor: int = 2,\n        feed_forward_dropout_p: float = 0.0,\n        attention_dropout_p: float = 0.1,\n        conv_dropout_p: float = 0.1,\n        conv_kernel_size: int = 31,\n        use_bn: bool = True,\n    ):\n        super(SqueezeformerBlock, self).__init__()\n\n        self.scale_mhsa, self.bias_mhsa = make_scale(encoder_dim)\n        self.scale_ff_mhsa, self.bias_ff_mhsa = make_scale(encoder_dim)\n        self.scale_conv, self.bias_conv = make_scale(encoder_dim)\n        self.scale_ff_conv, self.bias_ff_conv = make_scale(encoder_dim)\n\n        self.mhsa = MultiHeadedSelfAttentionModule(\n            d_model=encoder_dim,\n            num_heads=num_attention_heads,\n            dropout_p=attention_dropout_p,\n        )\n        self.ln_mhsa = nn.LayerNorm(encoder_dim)\n        self.ff_mhsa = FeedForwardModule(\n            encoder_dim=encoder_dim,\n            expansion_factor=feed_forward_expansion_factor,\n            dropout_p=feed_forward_dropout_p,\n        )\n        self.ln_ff_mhsa = nn.LayerNorm(encoder_dim)\n        self.conv = ConvModule(\n            in_channels=encoder_dim,\n            kernel_size=conv_kernel_size,\n            expansion_factor=conv_expansion_factor,\n            dropout_p=conv_dropout_p,\n            use_bn=use_bn,\n        )\n        self.ln_conv = nn.LayerNorm(encoder_dim)\n        self.ff_conv = FeedForwardModule(\n            encoder_dim=encoder_dim,\n            expansion_factor=feed_forward_expansion_factor,\n            dropout_p=feed_forward_dropout_p,\n        )\n        self.ln_ff_conv = nn.LayerNorm(encoder_dim)\n\n    def forward(self, x, mask):\n        mask_pad = (mask).long().bool().unsqueeze(1)\n        mask_pad = ~(mask_pad.permute(0, 2, 1) * mask_pad)\n        mask_flat = mask.view(-1).bool()\n        bs, slen, nfeats = x.shape\n\n        # MHSA\n        residual = x\n        x = x * self.scale_mhsa + self.bias_mhsa\n        x = residual + self.mhsa(x, mask_pad)\n\n        # Skip padding for layer norm\n        x_skip = x.reshape(-1, x.shape[-1])\n        x = x_skip[mask_flat].unsqueeze(0)\n        x = self.ln_mhsa(x)\n\n        # FFN after MHSA\n        residual = x\n        x = x * self.scale_ff_mhsa + self.bias_ff_mhsa\n        x = residual + self.ff_mhsa(x)\n        x = self.ln_ff_mhsa(x)\n\n        # Restore shape\n        x_skip[mask_flat] = x[0].to(dtype=x_skip.dtype)\n        x = x_skip.reshape(bs, slen, nfeats)\n\n        # Conv\n        residual = x\n        x = x * self.scale_conv + self.bias_conv\n        x = residual + self.conv(x, mask_pad=mask.bool().unsqueeze(1))\n\n        # Skip padding for layer norm\n        x_skip = x.reshape(-1, x.shape[-1])\n        x = x_skip[mask_flat].unsqueeze(0)\n        x = self.ln_conv(x)\n\n        # FFN after Conv\n        residual = x\n        x = x * self.scale_ff_conv + self.bias_ff_conv\n        x = residual + self.ff_conv(x)\n        x = self.ln_ff_conv(x)\n\n        # Restore shape\n        x_skip[mask_flat] = x[0].to(dtype=x_skip.dtype)\n        x = x_skip.reshape(bs, slen, nfeats)\n\n        return x\n\nclass SqueezeformerEncoder(nn.Module):\n    \"\"\"\n    Stack of Squeezeformer blocks\n    \"\"\"\n    def __init__(\n        self,\n        input_dim: int = 80,\n        encoder_dim: int = 512,\n        num_layers: int = 16,\n        num_attention_heads: int = 8,\n        feed_forward_expansion_factor: int = 4,\n        conv_expansion_factor: int = 2,\n        input_dropout_p: float = 0.0,\n        feed_forward_dropout_p: float = 0.0,\n        attention_dropout_p: float = 0.0,\n        conv_dropout_p: float = 0.0,\n        conv_kernel_size: int = 31,\n        use_bn: bool = True,\n    ):\n        super(SqueezeformerEncoder, self).__init__()\n        self.num_layers = num_layers\n\n        self.blocks = nn.ModuleList()\n        for idx in range(num_layers):\n            self.blocks.append(\n                SqueezeformerBlock(\n                    encoder_dim=encoder_dim,\n                    num_attention_heads=num_attention_heads,\n                    feed_forward_expansion_factor=feed_forward_expansion_factor,\n                    conv_expansion_factor=conv_expansion_factor,\n                    feed_forward_dropout_p=feed_forward_dropout_p,\n                    attention_dropout_p=attention_dropout_p,\n                    conv_dropout_p=conv_dropout_p,\n                    conv_kernel_size=conv_kernel_size,\n                    use_bn=use_bn,\n                )\n            )\n\n    def forward(self, x: Tensor, mask: Tensor):\n        for idx, block in enumerate(self.blocks):\n            x = block(x, mask)\n        return x\n\nfrom timm.layers.norm_act import BatchNormAct2d\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch_geometric.nn import TransformerConv, LayerNorm\n\nclass SpatialMouseGNN(nn.Module):\n    def __init__(self, in_channels, out_channels, num_mice=4, heads=4, dropout=0.1):\n        super().__init__()\n        self.num_mice = num_mice\n        # 1. Project Raw Features (Crucial Step)\n        # We need to turn coordinates into \"semantic\" embeddings before Attention\n        self.embedding = nn.Linear(in_channels, out_channels)\n        \n        # 2. TransformerConv Layers\n        # We use beta=True for a gating mechanism (improves deep GNNs)\n        self.conv1 = TransformerConv(out_channels, out_channels // heads, heads=heads, \n                                     dropout=dropout, beta=True)\n        self.norm1 = LayerNorm(out_channels) # LayerNorm is often more stable for Transformers\n        \n        self.conv2 = TransformerConv(out_channels, out_channels // heads, heads=heads, \n                                     dropout=dropout, beta=True)\n        self.norm2 = LayerNorm(out_channels)\n        \n        self.dropout = nn.Dropout(dropout)\n\n    def _get_fully_connected_edge_index(self, batch_size_time, device):\n        \"\"\"\n        Creates edges so every mouse connects to every other mouse within the same frame.\n        \"\"\"\n        # Base edges for one frame of 4 mice (0-3)\n        # 0->1, 0->2, 0->3, 1->0, ...\n        base_src = torch.tensor([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], device=device)\n        base_dst = torch.tensor([1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2], device=device)\n        # If you want self-loops (mice look at themselves), add them here, \n        # but TransformerConv usually adds them automatically or handles them separately.\n\n        # Repeat for the whole batch\n        # We need to offset indices by num_mice (4) for each subsequent frame\n        offsets = torch.arange(batch_size_time, device=device) * self.num_mice\n        \n        # Broadcasting to create the massive edge list\n        # shape: [Edges_Per_Frame, Batch_Time]\n        src = base_src.unsqueeze(1) + offsets.unsqueeze(0)\n        dst = base_dst.unsqueeze(1) + offsets.unsqueeze(0)\n        \n        edge_index = torch.stack([src.flatten(), dst.flatten()], dim=0)\n        return edge_index\n\n    def forward(self, x):\n        \"\"\"\n        Input: [Batch, Frames, Mice, Features]\n        Output: [Batch, Frames, Mice, Out_Channels]\n        \"\"\"\n        B, T, M, D = x.shape\n        \n        # Flatten Batch and Time: treat every frame as an independent graph\n        # New shape: [Total_Graphs * Mice, Features]\n        x_flat = x.reshape(B * T * M, D)\n        \n        # 1. Linear Projection\n        x_emb = F.relu(self.embedding(x_flat))\n        \n        # 2. Create Edges on the fly\n        # (Optimized: we only calculate this once per forward pass structure)\n        edge_index = self._get_fully_connected_edge_index(B * T, x.device)\n        \n        # 3. GNN Layer 1\n        # TransformerConv expects [Num_Nodes, Dim]\n        h = self.conv1(x_emb, edge_index)\n        h = self.norm1(h)\n        h = F.relu(h)\n        h = self.dropout(h)\n        \n        # 4. GNN Layer 2\n        h = self.conv2(h, edge_index)\n        h = self.norm2(h)\n        h = F.relu(h)\n        \n        # 5. Reshape back (NO POOLING)\n        # We want to keep the mice separate for the Squeezeformer\n        return h.reshape(B, T, M, -1)\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass RelationAugmentedHead(nn.Module):\n    def __init__(self, encoder_dim, num_classes=39, dropout=0.0):\n        super().__init__()\n        \n        # --- 1. The Relation Network Module (Global Context) ---\n        # g_theta: Processes every pair to find relationships\n        # Input: [Mouse_A + Mouse_B] -> 2 * encoder_dim\n        self.g_theta = nn.Sequential(\n            nn.Linear(encoder_dim * 2, encoder_dim),\n            nn.ReLU(),\n            nn.Linear(encoder_dim, encoder_dim // 2),\n            nn.ReLU()\n        )\n        \n        # f_phi: Processes the SUM of all relations (The \"Cage Vibe\")\n        self.f_phi = nn.Sequential(\n            nn.Linear(encoder_dim // 2, encoder_dim // 2),\n            nn.ReLU(),\n            nn.Dropout(dropout)\n        )\n        \n        # --- 2. The Pairwise Classifier ---\n        # Input: [Mouse_A + Mouse_B + Global_Context]\n        self.in_features = (encoder_dim * 2) + (encoder_dim // 2)\n        \n        self.classifier = nn.Sequential(\n            nn.Linear(self.in_features, encoder_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(encoder_dim, num_classes)\n        )\n\n    def forward(self, x):\n        \"\"\"\n        Args:\n            x: [Batch, Time, Mice, Dim] - Output from Squeezeformer/GNN\n        \"\"\"\n        B, T, M, D = x.shape\n        \n        # --- Step A: Form All Pairs (Same as before) ---\n        # [B, T, M, M, D]\n        m1 = x.unsqueeze(3).expand(-1, -1, -1, M, -1)\n        m2 = x.unsqueeze(2).expand(-1, -1, M, -1, -1)\n        \n        # The specific pairs we want to classify\n        # Shape: [B, T, M, M, 2*D]\n        specific_pairs = torch.cat([m1, m2], dim=-1)\n        \n        # --- Step B: The Relation Network (Global \"Thinking\") ---\n        # 1. Apply g_theta to ALL pairs\n        # Shape: [B, T, M, M, D/2]\n        relations = self.g_theta(specific_pairs)\n        \n        # 2. Sum over all pairs (M*M) to get ONE vector per frame\n        # This represents the \"Global Social State\" of the cage at time T\n        # Shape: [B, T, D/2]\n        global_sum = relations.sum(dim=(2, 3)) \n        \n        # 3. Apply f_phi (Reasoning over the global state)\n        # Shape: [B, T, D/2]\n        global_context = self.f_phi(global_sum)\n        \n        # --- Step C: Combine Local + Global ---\n        # We need to broadcast the global context to every specific pair\n        # Global: [B, T, 1, 1, D/2]\n        global_context_expanded = global_context.unsqueeze(2).unsqueeze(3).expand(-1, -1, M, M, -1)\n        \n        # Concatenate: [Specific_Pair (2D) | Global_Context (0.5D)]\n        # Shape: [B, T, M, M, 2.5D]\n        combined_features = torch.cat([specific_pairs, global_context_expanded], dim=-1)\n        \n        # --- Step D: Final Classification ---\n        # Flatten pairs: [B, T, 16, 2.5D]\n        combined_flat = combined_features.view(B, T, M*M, -1)\n        \n        # Predict\n        logits = self.classifier(combined_flat)\n        \n        return logits\n\nclass Net(nn.Module):\n    \"\"\"\n    Squeezeformer model for MABe mouse behavior detection\n    \"\"\"\n    def __init__(self, cfg):\n        super(Net, self).__init__()\n        self.cfg = cfg\n        # Model dimensions\n        self.encoder_dim = cfg.encoder_config.encoder_dim\n        self.num_pairs = 16  # Number of mouse pairs\n        self.num_actions = 39  # Number of behavior classes (including no_action)\n\n        self.gnn = SpatialMouseGNN(\n            in_channels=cfg.per_mouse_feature_dim,\n            out_channels=cfg.encoder_config.encoder_dim,\n            num_mice=4,\n            heads=4,\n            dropout=cfg.encoder_config.input_dropout_p\n        )\n\n        self.encoder = SqueezeformerEncoder(\n            input_dim=self.encoder_dim,\n            encoder_dim=self.encoder_dim,\n            num_layers=cfg.encoder_config.num_layers,\n            num_attention_heads=cfg.encoder_config.num_attention_heads,\n            feed_forward_expansion_factor=cfg.encoder_config.feed_forward_expansion_factor,\n            conv_expansion_factor=cfg.encoder_config.conv_expansion_factor,\n            input_dropout_p=cfg.encoder_config.input_dropout_p,\n            feed_forward_dropout_p=cfg.encoder_config.feed_forward_dropout_p,\n            attention_dropout_p=cfg.encoder_config.attention_dropout_p,\n            conv_dropout_p=cfg.encoder_config.conv_dropout_p,\n            conv_kernel_size=cfg.encoder_config.conv_kernel_size,\n            use_bn = cfg.use_bn\n        )\n\n        self.classifier = RelationAugmentedHead(\n            encoder_dim=self.encoder_dim,\n            num_classes=self.num_actions,\n            dropout=cfg.encoder_config.feed_forward_dropout_p\n        )\n\n        if hasattr(cfg, 'class_weights') and cfg.class_weights is not None:\n            class_weights = torch.tensor(cfg.class_weights)\n        else:\n            pos_weight = cfg.pos_weight if hasattr(cfg, 'pos_weight') else 50.0 \n            class_weights = torch.full((self.num_actions,), pos_weight) # High weight for all\n            no_action_idx = cfg.action_id_map.get('no_action', -1)\n            if no_action_idx != -1:\n                class_weights[no_action_idx] = 1.0 # Low weight for no_action\n            else:\n                # Fallback if cfg is wrong, assume last class is no_action\n                class_weights[-1] = 1.0\n\n        self.register_buffer('class_weights', class_weights)\n\n        # Use CrossEntropyLoss for multi-class classification\n        self.loss_fn = nn.CrossEntropyLoss(weight=self.class_weights, reduction='none')\n        self.unweighted_ce_fn = nn.CrossEntropyLoss(reduction='none')\n\n        # Optional: Multi-class focal loss\n        self.use_focal_loss = cfg.use_focal_loss if hasattr(cfg, 'use_focal_loss') else False\n        self.focal_gamma = cfg.focal_gamma if hasattr(cfg, 'focal_gamma') else 2.0\n\n        # Training settings\n        self.return_logits = cfg.return_logits if hasattr(cfg, 'return_logits') else False\n\n        print(f'Model initialized with {count_parameters(self):,} trainable parameters')\n        # print(f'Loss: {\"Multi-class Focal\" if self.use_focal_loss else \"Weighted CrossEntropy\"}')\n        # print(f'Class weights: {self.class_weights.cpu().numpy()}')\n\n    def forward(self, batch):\n        mask = batch['input_mask'].long()  # (batch, seq_len)\n        mask_for_encoder = mask.repeat_interleave(4, dim=0)\n\n        if self.cfg.cnn_extractor:\n            x = batch['input_mice']\n            x = self.gnn(x)  # (batch, seq_len, num_mice, encoder_dim)\n            B, T, M, C = x.shape\n            x = x.permute(0,2,1,3).reshape(B * M, T, C)  # (batch*num_mice, seq_len, encoder_dim)\n\n\n        # Encode\n        x = self.encoder(x, mask_for_encoder)  # (batch, seq_len, encoder_dim)\n        x = x.reshape(B, M, T, C).permute(0,2,1,3)  # (batch, seq_len, num_mice, encoder_dim)\n\n        # Classify\n        logits = self.classifier(x)  # (batch, seq_len, num_pairs, num_actions)\n\n        # APPLY BEHAVIOR MASK if provided (for evaluation)\n        if 'behavior_mask' in batch and batch['behavior_mask'] is not None:\n            behavior_mask = batch['behavior_mask']  # (batch, num_pairs, num_actions)\n            # Expand to match logits shape\n            behavior_mask_expanded = behavior_mask.unsqueeze(1).expand_as(logits)\n            # For multi-class, we need to ensure at least one action is valid per pair\n            # Set masked positions to very negative value\n            logits = torch.where(behavior_mask_expanded.bool(), logits,\n                                torch.tensor(-1e10, dtype=logits.dtype, device=logits.device) if logits.dtype == torch.float32\n                                else torch.tensor(-1e4, dtype=logits.dtype, device=logits.device))\n\n        output = {}\n\n        # Calculate loss if labels provided\n        if 'labels' in batch and batch['labels'] is not None:\n            labels = batch['labels']\n\n            # Convert one-hot to class indices if needed\n            if labels.dim() == 4:  # One-hot encoded (batch, seq_len, num_pairs, num_actions)\n                # Convert to class indices\n                labels = torch.argmax(labels, dim=-1)  # (batch, seq_len, num_pairs)\n\n            # Reshape for loss calculation\n            batch_size, seq_len, num_pairs, num_actions = logits.shape\n            logits_flat = logits.reshape(-1, num_actions)  # (batch*seq*pairs, num_actions)\n            labels_flat = labels.reshape(-1)  # (batch*seq*pairs,)\n\n            # Calculate loss\n            if self.use_focal_loss:\n                # Multi-class focal loss\n                unweighted_ce = self.unweighted_ce_fn(logits_flat, labels_flat)\n                pt = torch.exp(-unweighted_ce)\n                alpha = self.class_weights[labels_flat]  # Gather alpha for true classes\n                focal_loss = alpha * (1 - pt) ** self.focal_gamma * unweighted_ce\n                loss = focal_loss\n            else:\n                # Standard weighted cross-entropy\n                loss = self.loss_fn(logits_flat, labels_flat)\n\n            # Reshape loss back\n            loss = loss.view(batch_size, seq_len, num_pairs)\n\n            mask_expanded = mask.unsqueeze(-1).expand_as(loss)\n            loss = loss * mask_expanded\n\n            # Average over valid positions\n            valid_positions = mask_expanded.sum()\n            if valid_positions > 0:\n                loss = loss.sum() / valid_positions\n            else:\n                loss = loss.sum()\n            output['loss'] = loss\n\n        probs = torch.softmax(logits, dim=-1)  # Multi-class probabilities\n        output['predictions'] = probs\n        output['logits'] = logits\n\n        # Store additional info for evaluation\n        if 'video_id' in batch:\n            output['video_id'] = batch['video_id']\n        if 'start_frame' in batch:\n            output['start_frame'] = batch['start_frame']\n\n        return output","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T13:06:49.009475Z","iopub.execute_input":"2025-12-18T13:06:49.009802Z","iopub.status.idle":"2025-12-18T13:06:49.033536Z","shell.execute_reply.started":"2025-12-18T13:06:49.009784Z","shell.execute_reply":"2025-12-18T13:06:49.032907Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### MODEL 242","metadata":{}},{"cell_type":"code","source":"%%writefile model_242.py\n\nimport torch\nfrom torch.nn import functional as F\nfrom torch import nn\nfrom typing import Tuple, Union, Optional\nfrom torch import Tensor\nimport math\nimport numpy as np\n\n\ndef count_parameters(model):\n    \"\"\"Count the number of trainable parameters in model\"\"\"\n    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n\nclass Swish(nn.Module):\n    \"\"\"Swish activation function\"\"\"\n    def __init__(self) -> None:\n        super(Swish, self).__init__()\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return inputs * inputs.sigmoid()\n\nclass GLU(nn.Module):\n    \"\"\"Gated Linear Unit activation\"\"\"\n    def __init__(self, dim: int) -> None:\n        super(GLU, self).__init__()\n        self.dim = dim\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        outputs, gate = inputs.chunk(2, dim=self.dim)\n        return outputs * gate.sigmoid()\n\nclass FeedForwardModule(nn.Module):\n    \"\"\"\n    Feed Forward Module with pre-norm residual units\n    \"\"\"\n    def __init__(\n        self,\n        encoder_dim: int = 512,\n        expansion_factor: int = 4,\n        dropout_p: float = 0.0,\n    ) -> None:\n        super(FeedForwardModule, self).__init__()\n\n        self.ffn1 = nn.Linear(encoder_dim, encoder_dim * expansion_factor, bias=True)\n        self.act = Swish()\n        self.do1 = nn.Dropout(p=dropout_p)\n        self.ffn2 = nn.Linear(encoder_dim * expansion_factor, encoder_dim, bias=True)\n        self.do2 = nn.Dropout(p=dropout_p)\n\n    def forward(self, x):\n        x = self.ffn1(x)\n        x = self.act(x)\n        x = self.do1(x)\n        x = self.ffn2(x)\n        x = self.do2(x)\n        return x\n\nclass RelPositionalEncoding(nn.Module):\n    \"\"\"\n    Relative positional encoding module for handling variable sequence lengths\n    \"\"\"\n    def __init__(self, d_model: int = 512, max_len: int = 5000) -> None:\n        super(RelPositionalEncoding, self).__init__()\n        self.d_model = d_model\n        self.pe = None\n        self.extend_pe(torch.tensor(0.0).expand(1, max_len))\n\n    def extend_pe(self, x):\n        if self.pe is not None:\n            if self.pe.size(1) >= x.size(1) * 2 - 1:\n                if self.pe.dtype != x.dtype or self.pe.device != x.device:\n                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)\n                return\n\n        pe_positive = torch.zeros(x.size(1), self.d_model)\n        pe_negative = torch.zeros(x.size(1), self.d_model)\n        position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)\n        div_term = torch.exp(\n            torch.arange(0, self.d_model, 2, dtype=torch.float32) * -(math.log(10000.0) / self.d_model)\n        )\n        pe_positive[:, 0::2] = torch.sin(position * div_term)\n        pe_positive[:, 1::2] = torch.cos(position * div_term)\n        pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)\n        pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)\n\n        pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)\n        pe_negative = pe_negative[1:].unsqueeze(0)\n        pe = torch.cat([pe_positive, pe_negative], dim=1)\n        self.pe = pe.to(device=x.device, dtype=x.dtype)\n\n    def forward(self, x: torch.Tensor):\n        self.extend_pe(x)\n        pos_emb = self.pe[\n            :,\n            self.pe.size(1) // 2 - x.size(1) + 1 : self.pe.size(1) // 2 + x.size(1),\n        ]\n        return pos_emb\n\nclass RelativeMultiHeadAttention(nn.Module):\n    \"\"\"\n    Multi-head attention with relative positional encoding\n    \"\"\"\n    def __init__(\n        self,\n        d_model: int = 512,\n        num_heads: int = 16,\n        dropout_p: float = 0.0,\n    ):\n        super(RelativeMultiHeadAttention, self).__init__()\n        assert d_model % num_heads == 0, \"d_model % num_heads should be zero.\"\n        self.d_model = d_model\n        self.d_head = int(d_model / num_heads)\n        self.num_heads = num_heads\n        self.sqrt_dim = math.sqrt(self.d_head)\n\n        self.query_proj = nn.Linear(d_model, d_model)\n        self.key_proj = nn.Linear(d_model, d_model)\n        self.value_proj = nn.Linear(d_model, d_model)\n        self.pos_proj = nn.Linear(d_model, d_model, bias=False)\n\n        self.dropout = nn.Dropout(p=dropout_p)\n        self.u_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))\n        self.v_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))\n        torch.nn.init.xavier_uniform_(self.u_bias)\n        torch.nn.init.xavier_uniform_(self.v_bias)\n\n        self.out_proj = nn.Linear(d_model, d_model)\n\n    def forward(\n        self,\n        query: Tensor,\n        key: Tensor,\n        value: Tensor,\n        pos_embedding: Tensor,\n        mask: Optional[Tensor] = None,\n    ) -> Tensor:\n        batch_size = value.size(0)\n\n        query = self.query_proj(query).view(batch_size, -1, self.num_heads, self.d_head)\n        key = self.key_proj(key).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)\n        value = self.value_proj(value).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)\n        pos_embedding = self.pos_proj(pos_embedding).view(batch_size, -1, self.num_heads, self.d_head)\n\n        content_score = torch.matmul((query + self.u_bias).transpose(1, 2), key.transpose(2, 3))\n        pos_score = torch.matmul((query + self.v_bias).transpose(1, 2), pos_embedding.permute(0, 2, 3, 1))\n        pos_score = self._relative_shift(pos_score)\n\n        score = (content_score + pos_score) / self.sqrt_dim\n\n        if mask is not None:\n            mask = mask.unsqueeze(1)\n            score.masked_fill_(mask, -1e9) if score.dtype == torch.float32 else score.masked_fill_(mask, -1e4)\n\n        attn = F.softmax(score, -1)\n        attn = self.dropout(attn)\n\n        context = torch.matmul(attn, value).transpose(1, 2)\n        context = context.contiguous().view(batch_size, -1, self.d_model)\n\n        return self.out_proj(context)\n\n    def _relative_shift(self, pos_score: Tensor) -> Tensor:\n        batch_size, num_heads, seq_length1, seq_length2 = pos_score.size()\n        zeros = pos_score.new_zeros(batch_size, num_heads, seq_length1, 1)\n        padded_pos_score = torch.cat([zeros, pos_score], dim=-1)\n\n        padded_pos_score = padded_pos_score.view(batch_size, num_heads, seq_length2 + 1, seq_length1)\n        pos_score = padded_pos_score[:, :, 1:].view_as(pos_score)[:, :, :, : seq_length2 // 2 + 1]\n\n        return pos_score\n\nclass MultiHeadedSelfAttentionModule(nn.Module):\n    \"\"\"\n    Self-attention module with relative positional encoding\n    \"\"\"\n    def __init__(self, d_model: int, num_heads: int, dropout_p: float = 0.0):\n        super(MultiHeadedSelfAttentionModule, self).__init__()\n        self.positional_encoding = RelPositionalEncoding(d_model)\n        self.attention = RelativeMultiHeadAttention(d_model, num_heads, dropout_p)\n        self.dropout = nn.Dropout(p=dropout_p)\n\n    def forward(self, inputs: Tensor, mask: Optional[Tensor] = None):\n        batch_size = inputs.size(0)\n        pos_embedding = self.positional_encoding(inputs)\n        pos_embedding = pos_embedding.repeat(batch_size, 1, 1)\n\n        outputs = self.attention(inputs, inputs, inputs, pos_embedding=pos_embedding, mask=mask)\n        return self.dropout(outputs)\n\nclass DepthwiseConv1d(nn.Module):\n    \"\"\"Depthwise 1D convolution\"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        kernel_size: int,\n        stride: int = 1,\n        padding: int = 0,\n        bias: bool = False,\n    ) -> None:\n        super(DepthwiseConv1d, self).__init__()\n        assert out_channels % in_channels == 0, \"out_channels should be constant multiple of in_channels\"\n        self.conv = nn.Conv1d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=kernel_size,\n            groups=in_channels,\n            stride=stride,\n            padding=padding,\n            bias=bias,\n        )\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return self.conv(inputs)\n\nclass PointwiseConv1d(nn.Module):\n    \"\"\"Pointwise 1D convolution (kernel size = 1)\"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        stride: int = 1,\n        padding: int = 0,\n        bias: bool = True,\n    ) -> None:\n        super(PointwiseConv1d, self).__init__()\n        self.conv = nn.Conv1d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=1,\n            stride=stride,\n            padding=padding,\n            bias=bias,\n        )\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return self.conv(inputs)\n\nclass ConvModule(nn.Module):\n    \"\"\"\n    Convolution module with pointwise conv -> GLU -> depthwise conv -> normalization -> activation\n    \"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        kernel_size: int = 31,\n        expansion_factor: int = 2,\n        dropout_p: float = 0.0,\n        use_bn: bool = True,\n    ) -> None:\n        super(ConvModule, self).__init__()\n        assert (kernel_size - 1) % 2 == 0, \"kernel_size should be a odd number for 'SAME' padding\"\n        assert expansion_factor == 2, \"Currently, Only Supports expansion_factor 2\"\n\n        self.pw_conv_1 = PointwiseConv1d(in_channels, in_channels * expansion_factor, stride=1, padding=0, bias=True)\n        self.act1 = GLU(dim=1)\n        self.dw_conv = DepthwiseConv1d(in_channels, in_channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2)\n        self.bn = nn.BatchNorm1d(in_channels)\n        self.inorm = nn.InstanceNorm1d(in_channels, affine=True)\n        self.act2 = Swish()\n        self.pw_conv_2 = PointwiseConv1d(in_channels, in_channels, stride=1, padding=0, bias=True)\n        self.do = nn.Dropout(p=dropout_p)\n        self.use_bn = use_bn\n\n    def forward(self, x, mask_pad):\n        # Transpose for conv operations [B, T, C]\n        x = x.transpose(1, 2)\n        if mask_pad.size(2) > 0:  # time > 0\n            x = x.masked_fill(~mask_pad, 0.0)\n\n        x = self.pw_conv_1(x)\n        x = self.act1(x)\n        x = self.dw_conv(x)\n\n        if self.use_bn:\n            # Apply batch norm only to non-padded positions\n            x_bn = x.permute(0,2,1).reshape(-1, x.shape[1])\n            mask_bn = mask_pad.view(-1)\n            x_bn[mask_bn] = self.bn(x_bn[mask_bn])\n            x = x_bn.view(x.permute(0,2,1).shape).permute(0,2,1)\n        else:    \n            x = self.inorm(x)\n\n        x = self.act2(x)\n        x = self.pw_conv_2(x)\n        x = self.do(x)\n\n        # Mask batch padding again\n        if mask_pad.size(2) > 0:  # time > 0\n            x = x.masked_fill(~mask_pad, 0.0)\n        x = x.transpose(1, 2)\n        return x\n\ndef make_scale(encoder_dim):\n    \"\"\"Create learnable scale and bias parameters\"\"\"\n    scale = torch.nn.Parameter(torch.tensor([1.] * encoder_dim)[None, None, :])\n    bias = torch.nn.Parameter(torch.tensor([0.] * encoder_dim)[None, None, :])\n    return scale, bias\n\nclass SqueezeformerBlock(nn.Module):\n    \"\"\"\n    Squeezeformer block: MHSA -> FFN -> Conv -> FFN with residual connections\n    \"\"\"\n    def __init__(\n        self,\n        encoder_dim: int = 512,\n        num_attention_heads: int = 8,\n        feed_forward_expansion_factor: int = 4,\n        conv_expansion_factor: int = 2,\n        feed_forward_dropout_p: float = 0.0,\n        attention_dropout_p: float = 0.1,\n        conv_dropout_p: float = 0.1,\n        conv_kernel_size: int = 31,\n        use_bn: bool = True,\n    ):\n        super(SqueezeformerBlock, self).__init__()\n\n        self.scale_mhsa, self.bias_mhsa = make_scale(encoder_dim)\n        self.scale_ff_mhsa, self.bias_ff_mhsa = make_scale(encoder_dim)\n        self.scale_conv, self.bias_conv = make_scale(encoder_dim)\n        self.scale_ff_conv, self.bias_ff_conv = make_scale(encoder_dim)\n\n        self.mhsa = MultiHeadedSelfAttentionModule(\n            d_model=encoder_dim,\n            num_heads=num_attention_heads,\n            dropout_p=attention_dropout_p,\n        )\n        self.ln_mhsa = nn.LayerNorm(encoder_dim)\n        self.ff_mhsa = FeedForwardModule(\n            encoder_dim=encoder_dim,\n            expansion_factor=feed_forward_expansion_factor,\n            dropout_p=feed_forward_dropout_p,\n        )\n        self.ln_ff_mhsa = nn.LayerNorm(encoder_dim)\n        self.conv = ConvModule(\n            in_channels=encoder_dim,\n            kernel_size=conv_kernel_size,\n            expansion_factor=conv_expansion_factor,\n            dropout_p=conv_dropout_p,\n            use_bn=use_bn,\n        )\n        self.ln_conv = nn.LayerNorm(encoder_dim)\n        self.ff_conv = FeedForwardModule(\n            encoder_dim=encoder_dim,\n            expansion_factor=feed_forward_expansion_factor,\n            dropout_p=feed_forward_dropout_p,\n        )\n        self.ln_ff_conv = nn.LayerNorm(encoder_dim)\n\n    def forward(self, x, mask):\n        mask_pad = (mask).long().bool().unsqueeze(1)\n        mask_pad = ~(mask_pad.permute(0, 2, 1) * mask_pad)\n        mask_flat = mask.view(-1).bool()\n        bs, slen, nfeats = x.shape\n\n        # MHSA\n        residual = x\n        x = x * self.scale_mhsa + self.bias_mhsa\n        x = residual + self.mhsa(x, mask_pad)\n\n        # Skip padding for layer norm\n        x_skip = x.view(-1, x.shape[-1])\n        x = x_skip[mask_flat].unsqueeze(0)\n        x = self.ln_mhsa(x)\n\n        # FFN after MHSA\n        residual = x\n        x = x * self.scale_ff_mhsa + self.bias_ff_mhsa\n        x = residual + self.ff_mhsa(x)\n        x = self.ln_ff_mhsa(x)\n\n        # Restore shape\n        x_skip[mask_flat] = x[0].to(dtype=x_skip.dtype)\n        x = x_skip.view(bs, slen, nfeats)\n\n        # Conv\n        residual = x\n        x = x * self.scale_conv + self.bias_conv\n        x = residual + self.conv(x, mask_pad=mask.bool().unsqueeze(1))\n\n        # Skip padding for layer norm\n        x_skip = x.view(-1, x.shape[-1])\n        x = x_skip[mask_flat].unsqueeze(0)\n        x = self.ln_conv(x)\n\n        # FFN after Conv\n        residual = x\n        x = x * self.scale_ff_conv + self.bias_ff_conv\n        x = residual + self.ff_conv(x)\n        x = self.ln_ff_conv(x)\n\n        # Restore shape\n        x_skip[mask_flat] = x[0].to(dtype=x_skip.dtype)\n        x = x_skip.view(bs, slen, nfeats)\n\n        return x\n\nclass SqueezeformerEncoder(nn.Module):\n    \"\"\"\n    Stack of Squeezeformer blocks\n    \"\"\"\n    def __init__(\n        self,\n        input_dim: int = 80,\n        encoder_dim: int = 512,\n        num_layers: int = 16,\n        num_attention_heads: int = 8,\n        feed_forward_expansion_factor: int = 4,\n        conv_expansion_factor: int = 2,\n        input_dropout_p: float = 0.0,\n        feed_forward_dropout_p: float = 0.0,\n        attention_dropout_p: float = 0.0,\n        conv_dropout_p: float = 0.0,\n        conv_kernel_size: int = 31,\n        use_bn: bool = True,\n    ):\n        super(SqueezeformerEncoder, self).__init__()\n        self.num_layers = num_layers\n\n        self.blocks = nn.ModuleList()\n        for idx in range(num_layers):\n            self.blocks.append(\n                SqueezeformerBlock(\n                    encoder_dim=encoder_dim,\n                    num_attention_heads=num_attention_heads,\n                    feed_forward_expansion_factor=feed_forward_expansion_factor,\n                    conv_expansion_factor=conv_expansion_factor,\n                    feed_forward_dropout_p=feed_forward_dropout_p,\n                    attention_dropout_p=attention_dropout_p,\n                    conv_dropout_p=conv_dropout_p,\n                    conv_kernel_size=conv_kernel_size,\n                    use_bn=use_bn,\n                )\n            )\n\n    def forward(self, x: Tensor, mask: Tensor):\n        for idx, block in enumerate(self.blocks):\n            x = block(x, mask)\n        return x\n\nfrom timm.layers.norm_act import BatchNormAct2d\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch_geometric.nn import TransformerConv, LayerNorm\n\nclass SpatialMouseGNN(nn.Module):\n    def __init__(self, in_channels, out_channels, edge_dim=6, num_mice=4, heads=4, dropout=0.1):\n        super().__init__()\n        self.num_mice = num_mice\n        \n        # 1. Project Raw Node Features\n        self.embedding = nn.Linear(in_channels, out_channels)\n        \n        # 2. Project Raw Edge Features (New)\n        # Projects your 6 social features to match the GNN's internal dimension\n        self.edge_embedding = nn.Linear(edge_dim, edge_dim)\n        \n        # 3. TransformerConv Layers\n        # Added edge_dim argument so it knows what to expect\n        self.conv1 = TransformerConv(out_channels, out_channels // heads, heads=heads, \n                                     dropout=dropout, beta=True, edge_dim=edge_dim)\n        self.norm1 = LayerNorm(out_channels)\n        \n        self.conv2 = TransformerConv(out_channels, out_channels // heads, heads=heads, \n                                     dropout=dropout, beta=True, edge_dim=edge_dim)\n        self.norm2 = LayerNorm(out_channels)\n        \n        self.dropout = nn.Dropout(dropout)\n\n    def _get_fully_connected_edge_index(self, batch_size_time, device):\n        \"\"\"\n        Creates edges so every mouse connects to every other mouse within the same frame.\n        \"\"\"\n        # Base edges for one frame of 4 mice (0-3)\n        # 0->1, 0->2, 0->3, 1->0, ...\n        base_src = torch.tensor([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], device=device)\n        base_dst = torch.tensor([1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2], device=device)\n        \n        # Offsets for the whole batch\n        offsets = torch.arange(batch_size_time, device=device) * self.num_mice\n        \n        # Broadcasting to create [Batch_Time, Edges_Per_Frame]\n        # We use (BT, 1) + (1, 12) to get (BT, 12)\n        # This ensures edges are ordered: [Graph0_Edges, Graph1_Edges, ...]\n        src = offsets.unsqueeze(1) + base_src.unsqueeze(0)\n        dst = offsets.unsqueeze(1) + base_dst.unsqueeze(0)\n        \n        edge_index = torch.stack([src.flatten(), dst.flatten()], dim=0)\n        return edge_index\n\n    def forward(self, x, x_edges):\n        \"\"\"\n        x: [Batch, Frames, Mice, Node_Features]\n        x_edges: [Batch, Frames, Mice, Mice, Edge_Features]\n        \"\"\"\n        B, T, M, D_n = x.shape\n        x_edges = x_edges[:, :, :, :, :6]  # [B, T, M, M, Edge_Features]\n        _, _, _, _, D_e = x_edges.shape\n        \n        # 1. Flatten Nodes: [Total_Mice, Dim]\n        x_flat = x.view(B * T * M, D_n)\n        x_emb = F.relu(self.embedding(x_flat))\n        \n        # 2. Flatten Edge Attributes to match Edge Index\n        # Reshape to [Batch*Time, M, M, Edge_Dim]\n        edge_feats_flat = x_edges.view(B * T, M, M, D_e)\n        \n        # We need to extract the specific 12 edges (0->1, 0->2...) per graph\n        # Base indices for one graph\n        base_src_idx = torch.tensor([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], device=x.device)\n        base_dst_idx = torch.tensor([1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2], device=x.device)\n        \n        # Select attributes for valid edges\n        # Shape: [Batch*Time, 12, Edge_Dim]\n        valid_edge_attrs = edge_feats_flat[:, base_src_idx, base_dst_idx, :]\n        \n        # Flatten to [Total_Edges, Edge_Dim]\n        # This matches the order of _get_fully_connected_edge_index\n        edge_attr = valid_edge_attrs.view(-1, D_e)\n        \n        # Project Edge Attributes\n        edge_attr = F.relu(self.edge_embedding(edge_attr))\n        \n        # 3. Create Edge Index\n        edge_index = self._get_fully_connected_edge_index(B * T, x.device)\n        \n        # 4. GNN Layers\n        h = self.conv1(x_emb, edge_index, edge_attr=edge_attr)\n        h = self.norm1(h)\n        h = F.relu(h)\n        h = self.dropout(h)\n        \n        h = self.conv2(h, edge_index, edge_attr=edge_attr)\n        h = self.norm2(h)\n        h = F.relu(h)\n        \n        # 5. Reshape back\n        return h.view(B, T, M, -1)\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass RelationAugmentedHead(nn.Module):\n    def __init__(self, encoder_dim, num_classes=39, dropout=0.0):\n        super().__init__()\n        \n        # --- 1. The Relation Network Module (Global Context) ---\n        # g_theta: Processes every pair to find relationships\n        # Input: [Mouse_A + Mouse_B] -> 2 * encoder_dim\n        self.g_theta = nn.Sequential(\n            nn.Linear(encoder_dim * 2, encoder_dim),\n            nn.ReLU(),\n            nn.Linear(encoder_dim, encoder_dim // 2),\n            nn.ReLU()\n        )\n        \n        # f_phi: Processes the SUM of all relations (The \"Cage Vibe\")\n        self.f_phi = nn.Sequential(\n            nn.Linear(encoder_dim // 2, encoder_dim // 2),\n            nn.ReLU(),\n            nn.Dropout(dropout)\n        )\n        \n        # --- 2. The Pairwise Classifier ---\n        # Input: [Mouse_A + Mouse_B + Global_Context]\n        self.in_features = (encoder_dim * 2) + (encoder_dim // 2)\n        \n        self.classifier = nn.Sequential(\n            nn.Linear(self.in_features, encoder_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(encoder_dim, num_classes)\n        )\n\n    def forward(self, x):\n        \"\"\"\n        Args:\n            x: [Batch, Time, Mice, Dim] - Output from Squeezeformer/GNN\n        \"\"\"\n        B, T, M, D = x.shape\n        \n        # --- Step A: Form All Pairs (Same as before) ---\n        # [B, T, M, M, D]\n        m1 = x.unsqueeze(3).expand(-1, -1, -1, M, -1)\n        m2 = x.unsqueeze(2).expand(-1, -1, M, -1, -1)\n        \n        # The specific pairs we want to classify\n        # Shape: [B, T, M, M, 2*D]\n        specific_pairs = torch.cat([m1, m2], dim=-1)\n        \n        # --- Step B: The Relation Network (Global \"Thinking\") ---\n        # 1. Apply g_theta to ALL pairs\n        # Shape: [B, T, M, M, D/2]\n        relations = self.g_theta(specific_pairs)\n        \n        # 2. Sum over all pairs (M*M) to get ONE vector per frame\n        # This represents the \"Global Social State\" of the cage at time T\n        # Shape: [B, T, D/2]\n        global_sum = relations.sum(dim=(2, 3)) \n        \n        # 3. Apply f_phi (Reasoning over the global state)\n        # Shape: [B, T, D/2]\n        global_context = self.f_phi(global_sum)\n        \n        # --- Step C: Combine Local + Global ---\n        # We need to broadcast the global context to every specific pair\n        # Global: [B, T, 1, 1, D/2]\n        global_context_expanded = global_context.unsqueeze(2).unsqueeze(3).expand(-1, -1, M, M, -1)\n        \n        # Concatenate: [Specific_Pair (2D) | Global_Context (0.5D)]\n        # Shape: [B, T, M, M, 2.5D]\n        combined_features = torch.cat([specific_pairs, global_context_expanded], dim=-1)\n        \n        # --- Step D: Final Classification ---\n        # Flatten pairs: [B, T, 16, 2.5D]\n        combined_flat = combined_features.view(B, T, M*M, -1)\n        \n        # Predict\n        logits = self.classifier(combined_flat)\n        \n        return logits\n\nclass Net(nn.Module):\n    \"\"\"\n    Squeezeformer model for MABe mouse behavior detection\n    \"\"\"\n    def __init__(self, cfg):\n        super(Net, self).__init__()\n        self.cfg = cfg\n        # Model dimensions\n        self.encoder_dim = cfg.encoder_config.encoder_dim\n        self.num_pairs = 16  # Number of mouse pairs\n        self.num_actions = 39  # Number of behavior classes (including no_action)\n\n        self.gnn = SpatialMouseGNN(\n            in_channels=cfg.per_mouse_feature_dim,\n            out_channels=cfg.encoder_config.encoder_dim,\n            num_mice=4,\n            heads=4,\n            dropout=cfg.encoder_config.input_dropout_p\n        )\n\n        self.encoder = SqueezeformerEncoder(\n            input_dim=self.encoder_dim,\n            encoder_dim=self.encoder_dim,\n            num_layers=cfg.encoder_config.num_layers,\n            num_attention_heads=cfg.encoder_config.num_attention_heads,\n            feed_forward_expansion_factor=cfg.encoder_config.feed_forward_expansion_factor,\n            conv_expansion_factor=cfg.encoder_config.conv_expansion_factor,\n            input_dropout_p=cfg.encoder_config.input_dropout_p,\n            feed_forward_dropout_p=cfg.encoder_config.feed_forward_dropout_p,\n            attention_dropout_p=cfg.encoder_config.attention_dropout_p,\n            conv_dropout_p=cfg.encoder_config.conv_dropout_p,\n            conv_kernel_size=cfg.encoder_config.conv_kernel_size,\n            use_bn = cfg.use_bn\n        )\n\n        self.classifier = RelationAugmentedHead(\n            encoder_dim=self.encoder_dim,\n            num_classes=self.num_actions,\n            dropout=cfg.encoder_config.feed_forward_dropout_p\n        )\n\n        if hasattr(cfg, 'class_weights') and cfg.class_weights is not None:\n            class_weights = torch.tensor(cfg.class_weights)\n        else:\n            pos_weight = cfg.pos_weight if hasattr(cfg, 'pos_weight') else 50.0 \n            class_weights = torch.full((self.num_actions,), pos_weight) # High weight for all\n            no_action_idx = cfg.action_id_map.get('no_action', -1)\n            if no_action_idx != -1:\n                class_weights[no_action_idx] = 1.0 # Low weight for no_action\n            else:\n                # Fallback if cfg is wrong, assume last class is no_action\n                class_weights[-1] = 1.0\n\n        self.register_buffer('class_weights', class_weights)\n\n        # Use CrossEntropyLoss for multi-class classification\n        self.loss_fn = nn.CrossEntropyLoss(weight=self.class_weights, reduction='none')\n        self.unweighted_ce_fn = nn.CrossEntropyLoss(reduction='none')\n\n        # Optional: Multi-class focal loss\n        self.use_focal_loss = cfg.use_focal_loss if hasattr(cfg, 'use_focal_loss') else False\n        self.focal_gamma = cfg.focal_gamma if hasattr(cfg, 'focal_gamma') else 2.0\n\n        # Training settings\n        self.return_logits = cfg.return_logits if hasattr(cfg, 'return_logits') else False\n\n        print(f'Model initialized with {count_parameters(self):,} trainable parameters')\n        # print(f'Loss: {\"Multi-class Focal\" if self.use_focal_loss else \"Weighted CrossEntropy\"}')\n        # print(f'Class weights: {self.class_weights.cpu().numpy()}')\n\n    def forward(self, batch):\n        mask = batch['input_mask'].long()  # (batch, seq_len)\n        mask_for_encoder = mask.repeat_interleave(4, dim=0)\n\n        if self.cfg.cnn_extractor:\n            x = batch['input_mice']\n            x_nodes = batch['node_feats']  # Shape [B, T, 4, Node_Dim]\n            x_edges = batch['edge_feats']  # Shape [B, T, 4, 4, Edge_Dim]\n            x = self.gnn(x_nodes, x_edges)  # (batch, seq_len, num_mice, encoder_dim)\n            B, T, M, C = x.shape\n            x = x.permute(0,2,1,3).reshape(B * M, T, C)  # (batch*num_mice, seq_len, encoder_dim)\n\n\n        # Encode\n        x = self.encoder(x, mask_for_encoder)  # (batch, seq_len, encoder_dim)\n        x = x.reshape(B, M, T, C).permute(0,2,1,3)  # (batch, seq_len, num_mice, encoder_dim)\n\n        # Classify\n        logits = self.classifier(x)  # (batch, seq_len, num_pairs, num_actions)\n\n        # APPLY BEHAVIOR MASK if provided (for evaluation)\n        if 'behavior_mask' in batch and batch['behavior_mask'] is not None:\n            behavior_mask = batch['behavior_mask']  # (batch, num_pairs, num_actions)\n            # Expand to match logits shape\n            behavior_mask_expanded = behavior_mask.unsqueeze(1).expand_as(logits)\n            # For multi-class, we need to ensure at least one action is valid per pair\n            # Set masked positions to very negative value\n            logits = torch.where(behavior_mask_expanded.bool(), logits,\n                                torch.tensor(-1e10, dtype=logits.dtype, device=logits.device) if logits.dtype == torch.float32\n                                else torch.tensor(-1e4, dtype=logits.dtype, device=logits.device))\n\n        output = {}\n\n        # Calculate loss if labels provided\n        if 'labels' in batch and batch['labels'] is not None:\n            labels = batch['labels']\n\n            # Convert one-hot to class indices if needed\n            if labels.dim() == 4:  # One-hot encoded (batch, seq_len, num_pairs, num_actions)\n                # Convert to class indices\n                labels = torch.argmax(labels, dim=-1)  # (batch, seq_len, num_pairs)\n\n            # Reshape for loss calculation\n            batch_size, seq_len, num_pairs, num_actions = logits.shape\n            logits_flat = logits.reshape(-1, num_actions)  # (batch*seq*pairs, num_actions)\n            labels_flat = labels.reshape(-1)  # (batch*seq*pairs,)\n\n            # Calculate loss\n            if self.use_focal_loss:\n                # Multi-class focal loss\n                unweighted_ce = self.unweighted_ce_fn(logits_flat, labels_flat)\n                pt = torch.exp(-unweighted_ce)\n                alpha = self.class_weights[labels_flat]  # Gather alpha for true classes\n                focal_loss = alpha * (1 - pt) ** self.focal_gamma * unweighted_ce\n                loss = focal_loss\n            else:\n                # Standard weighted cross-entropy\n                loss = self.loss_fn(logits_flat, labels_flat)\n\n            # Reshape loss back\n            loss = loss.view(batch_size, seq_len, num_pairs)\n\n            mask_expanded = mask.unsqueeze(-1).expand_as(loss)\n            loss = loss * mask_expanded\n\n            # Average over valid positions\n            valid_positions = mask_expanded.sum()\n            if valid_positions > 0:\n                loss = loss.sum() / valid_positions\n            else:\n                loss = loss.sum()\n            output['loss'] = loss\n\n        probs = torch.softmax(logits, dim=-1)  # Multi-class probabilities\n        output['predictions'] = probs\n        output['logits'] = logits\n\n        # Store additional info for evaluation\n        if 'video_id' in batch:\n            output['video_id'] = batch['video_id']\n        if 'start_frame' in batch:\n            output['start_frame'] = batch['start_frame']\n\n        return output","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-12-18T13:06:49.034434Z","iopub.execute_input":"2025-12-18T13:06:49.034746Z","iopub.status.idle":"2025-12-18T13:06:49.063851Z","shell.execute_reply.started":"2025-12-18T13:06:49.034727Z","shell.execute_reply":"2025-12-18T13:06:49.063027Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### MODEL 243","metadata":{}},{"cell_type":"code","source":"%%writefile model_243.py\n\n\nimport torch\nfrom torch.nn import functional as F\nfrom torch import nn\nfrom typing import Tuple, Union, Optional\nfrom torch import Tensor\nimport math\nimport numpy as np\n\n\ndef count_parameters(model):\n    \"\"\"Count the number of trainable parameters in model\"\"\"\n    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n\nclass Swish(nn.Module):\n    \"\"\"Swish activation function\"\"\"\n    def __init__(self) -> None:\n        super(Swish, self).__init__()\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return inputs * inputs.sigmoid()\n\nclass GLU(nn.Module):\n    \"\"\"Gated Linear Unit activation\"\"\"\n    def __init__(self, dim: int) -> None:\n        super(GLU, self).__init__()\n        self.dim = dim\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        outputs, gate = inputs.chunk(2, dim=self.dim)\n        return outputs * gate.sigmoid()\n\nclass FeedForwardModule(nn.Module):\n    \"\"\"\n    Feed Forward Module with pre-norm residual units\n    \"\"\"\n    def __init__(\n        self,\n        encoder_dim: int = 512,\n        expansion_factor: int = 4,\n        dropout_p: float = 0.0,\n    ) -> None:\n        super(FeedForwardModule, self).__init__()\n\n        self.ffn1 = nn.Linear(encoder_dim, encoder_dim * expansion_factor, bias=True)\n        self.act = Swish()\n        self.do1 = nn.Dropout(p=dropout_p)\n        self.ffn2 = nn.Linear(encoder_dim * expansion_factor, encoder_dim, bias=True)\n        self.do2 = nn.Dropout(p=dropout_p)\n\n    def forward(self, x):\n        x = self.ffn1(x)\n        x = self.act(x)\n        x = self.do1(x)\n        x = self.ffn2(x)\n        x = self.do2(x)\n        return x\n\nclass RelPositionalEncoding(nn.Module):\n    \"\"\"\n    Relative positional encoding module for handling variable sequence lengths\n    \"\"\"\n    def __init__(self, d_model: int = 512, max_len: int = 5000) -> None:\n        super(RelPositionalEncoding, self).__init__()\n        self.d_model = d_model\n        self.pe = None\n        self.extend_pe(torch.tensor(0.0).expand(1, max_len))\n\n    def extend_pe(self, x):\n        if self.pe is not None:\n            if self.pe.size(1) >= x.size(1) * 2 - 1:\n                if self.pe.dtype != x.dtype or self.pe.device != x.device:\n                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)\n                return\n\n        pe_positive = torch.zeros(x.size(1), self.d_model)\n        pe_negative = torch.zeros(x.size(1), self.d_model)\n        position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)\n        div_term = torch.exp(\n            torch.arange(0, self.d_model, 2, dtype=torch.float32) * -(math.log(10000.0) / self.d_model)\n        )\n        pe_positive[:, 0::2] = torch.sin(position * div_term)\n        pe_positive[:, 1::2] = torch.cos(position * div_term)\n        pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)\n        pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)\n\n        pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)\n        pe_negative = pe_negative[1:].unsqueeze(0)\n        pe = torch.cat([pe_positive, pe_negative], dim=1)\n        self.pe = pe.to(device=x.device, dtype=x.dtype)\n\n    def forward(self, x: torch.Tensor):\n        self.extend_pe(x)\n        pos_emb = self.pe[\n            :,\n            self.pe.size(1) // 2 - x.size(1) + 1 : self.pe.size(1) // 2 + x.size(1),\n        ]\n        return pos_emb\n\nclass RelativeMultiHeadAttention(nn.Module):\n    \"\"\"\n    Multi-head attention with relative positional encoding\n    \"\"\"\n    def __init__(\n        self,\n        d_model: int = 512,\n        num_heads: int = 16,\n        dropout_p: float = 0.0,\n    ):\n        super(RelativeMultiHeadAttention, self).__init__()\n        assert d_model % num_heads == 0, \"d_model % num_heads should be zero.\"\n        self.d_model = d_model\n        self.d_head = int(d_model / num_heads)\n        self.num_heads = num_heads\n        self.sqrt_dim = math.sqrt(self.d_head)\n\n        self.query_proj = nn.Linear(d_model, d_model)\n        self.key_proj = nn.Linear(d_model, d_model)\n        self.value_proj = nn.Linear(d_model, d_model)\n        self.pos_proj = nn.Linear(d_model, d_model, bias=False)\n\n        self.dropout = nn.Dropout(p=dropout_p)\n        self.u_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))\n        self.v_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))\n        torch.nn.init.xavier_uniform_(self.u_bias)\n        torch.nn.init.xavier_uniform_(self.v_bias)\n\n        self.out_proj = nn.Linear(d_model, d_model)\n\n    def forward(\n        self,\n        query: Tensor,\n        key: Tensor,\n        value: Tensor,\n        pos_embedding: Tensor,\n        mask: Optional[Tensor] = None,\n    ) -> Tensor:\n        batch_size = value.size(0)\n\n        query = self.query_proj(query).view(batch_size, -1, self.num_heads, self.d_head)\n        key = self.key_proj(key).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)\n        value = self.value_proj(value).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)\n        pos_embedding = self.pos_proj(pos_embedding).view(batch_size, -1, self.num_heads, self.d_head)\n\n        content_score = torch.matmul((query + self.u_bias).transpose(1, 2), key.transpose(2, 3))\n        pos_score = torch.matmul((query + self.v_bias).transpose(1, 2), pos_embedding.permute(0, 2, 3, 1))\n        pos_score = self._relative_shift(pos_score)\n\n        score = (content_score + pos_score) / self.sqrt_dim\n\n        if mask is not None:\n            mask = mask.unsqueeze(1)\n            score.masked_fill_(mask, -1e9) if score.dtype == torch.float32 else score.masked_fill_(mask, -1e4)\n\n        attn = F.softmax(score, -1)\n        attn = self.dropout(attn)\n\n        context = torch.matmul(attn, value).transpose(1, 2)\n        context = context.contiguous().view(batch_size, -1, self.d_model)\n\n        return self.out_proj(context)\n\n    def _relative_shift(self, pos_score: Tensor) -> Tensor:\n        batch_size, num_heads, seq_length1, seq_length2 = pos_score.size()\n        zeros = pos_score.new_zeros(batch_size, num_heads, seq_length1, 1)\n        padded_pos_score = torch.cat([zeros, pos_score], dim=-1)\n\n        padded_pos_score = padded_pos_score.view(batch_size, num_heads, seq_length2 + 1, seq_length1)\n        pos_score = padded_pos_score[:, :, 1:].view_as(pos_score)[:, :, :, : seq_length2 // 2 + 1]\n\n        return pos_score\n\nclass MultiHeadedSelfAttentionModule(nn.Module):\n    \"\"\"\n    Self-attention module with relative positional encoding\n    \"\"\"\n    def __init__(self, d_model: int, num_heads: int, dropout_p: float = 0.0):\n        super(MultiHeadedSelfAttentionModule, self).__init__()\n        self.positional_encoding = RelPositionalEncoding(d_model)\n        self.attention = RelativeMultiHeadAttention(d_model, num_heads, dropout_p)\n        self.dropout = nn.Dropout(p=dropout_p)\n\n    def forward(self, inputs: Tensor, mask: Optional[Tensor] = None):\n        batch_size = inputs.size(0)\n        pos_embedding = self.positional_encoding(inputs)\n        pos_embedding = pos_embedding.repeat(batch_size, 1, 1)\n\n        outputs = self.attention(inputs, inputs, inputs, pos_embedding=pos_embedding, mask=mask)\n        return self.dropout(outputs)\n\nclass DepthwiseConv1d(nn.Module):\n    \"\"\"Depthwise 1D convolution\"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        kernel_size: int,\n        stride: int = 1,\n        padding: int = 0,\n        bias: bool = False,\n    ) -> None:\n        super(DepthwiseConv1d, self).__init__()\n        assert out_channels % in_channels == 0, \"out_channels should be constant multiple of in_channels\"\n        self.conv = nn.Conv1d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=kernel_size,\n            groups=in_channels,\n            stride=stride,\n            padding=padding,\n            bias=bias,\n        )\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return self.conv(inputs)\n\nclass PointwiseConv1d(nn.Module):\n    \"\"\"Pointwise 1D convolution (kernel size = 1)\"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        stride: int = 1,\n        padding: int = 0,\n        bias: bool = True,\n    ) -> None:\n        super(PointwiseConv1d, self).__init__()\n        self.conv = nn.Conv1d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=1,\n            stride=stride,\n            padding=padding,\n            bias=bias,\n        )\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return self.conv(inputs)\n\nclass ConvModule(nn.Module):\n    \"\"\"\n    Convolution module with pointwise conv -> GLU -> depthwise conv -> normalization -> activation\n    \"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        kernel_size: int = 31,\n        expansion_factor: int = 2,\n        dropout_p: float = 0.0,\n        use_bn: bool = True,\n    ) -> None:\n        super(ConvModule, self).__init__()\n        assert (kernel_size - 1) % 2 == 0, \"kernel_size should be a odd number for 'SAME' padding\"\n        assert expansion_factor == 2, \"Currently, Only Supports expansion_factor 2\"\n\n        self.pw_conv_1 = PointwiseConv1d(in_channels, in_channels * expansion_factor, stride=1, padding=0, bias=True)\n        self.act1 = GLU(dim=1)\n        self.dw_conv = DepthwiseConv1d(in_channels, in_channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2)\n        self.bn = nn.BatchNorm1d(in_channels)\n        self.inorm = nn.InstanceNorm1d(in_channels, affine=True)\n        self.act2 = Swish()\n        self.pw_conv_2 = PointwiseConv1d(in_channels, in_channels, stride=1, padding=0, bias=True)\n        self.do = nn.Dropout(p=dropout_p)\n        self.use_bn = use_bn\n\n    def forward(self, x, mask_pad):\n        # Transpose for conv operations [B, T, C]\n        x = x.transpose(1, 2)\n        if mask_pad.size(2) > 0:  # time > 0\n            x = x.masked_fill(~mask_pad, 0.0)\n\n        x = self.pw_conv_1(x)\n        x = self.act1(x)\n        x = self.dw_conv(x)\n\n        if self.use_bn:\n            # Apply batch norm only to non-padded positions\n            x_bn = x.permute(0,2,1).reshape(-1, x.shape[1])\n            mask_bn = mask_pad.view(-1)\n            x_bn[mask_bn] = self.bn(x_bn[mask_bn])\n            x = x_bn.view(x.permute(0,2,1).shape).permute(0,2,1)\n        else:    \n            x = self.inorm(x)\n\n        x = self.act2(x)\n        x = self.pw_conv_2(x)\n        x = self.do(x)\n\n        # Mask batch padding again\n        if mask_pad.size(2) > 0:  # time > 0\n            x = x.masked_fill(~mask_pad, 0.0)\n        x = x.transpose(1, 2)\n        return x\n\ndef make_scale(encoder_dim):\n    \"\"\"Create learnable scale and bias parameters\"\"\"\n    scale = torch.nn.Parameter(torch.tensor([1.] * encoder_dim)[None, None, :])\n    bias = torch.nn.Parameter(torch.tensor([0.] * encoder_dim)[None, None, :])\n    return scale, bias\n\nclass SqueezeformerBlock(nn.Module):\n    \"\"\"\n    Squeezeformer block: MHSA -> FFN -> Conv -> FFN with residual connections\n    \"\"\"\n    def __init__(\n        self,\n        encoder_dim: int = 512,\n        num_attention_heads: int = 8,\n        feed_forward_expansion_factor: int = 4,\n        conv_expansion_factor: int = 2,\n        feed_forward_dropout_p: float = 0.0,\n        attention_dropout_p: float = 0.1,\n        conv_dropout_p: float = 0.1,\n        conv_kernel_size: int = 31,\n        use_bn: bool = True,\n    ):\n        super(SqueezeformerBlock, self).__init__()\n\n        self.scale_mhsa, self.bias_mhsa = make_scale(encoder_dim)\n        self.scale_ff_mhsa, self.bias_ff_mhsa = make_scale(encoder_dim)\n        self.scale_conv, self.bias_conv = make_scale(encoder_dim)\n        self.scale_ff_conv, self.bias_ff_conv = make_scale(encoder_dim)\n\n        self.mhsa = MultiHeadedSelfAttentionModule(\n            d_model=encoder_dim,\n            num_heads=num_attention_heads,\n            dropout_p=attention_dropout_p,\n        )\n        self.ln_mhsa = nn.LayerNorm(encoder_dim)\n        self.ff_mhsa = FeedForwardModule(\n            encoder_dim=encoder_dim,\n            expansion_factor=feed_forward_expansion_factor,\n            dropout_p=feed_forward_dropout_p,\n        )\n        self.ln_ff_mhsa = nn.LayerNorm(encoder_dim)\n        self.conv = ConvModule(\n            in_channels=encoder_dim,\n            kernel_size=conv_kernel_size,\n            expansion_factor=conv_expansion_factor,\n            dropout_p=conv_dropout_p,\n            use_bn=use_bn,\n        )\n        self.ln_conv = nn.LayerNorm(encoder_dim)\n        self.ff_conv = FeedForwardModule(\n            encoder_dim=encoder_dim,\n            expansion_factor=feed_forward_expansion_factor,\n            dropout_p=feed_forward_dropout_p,\n        )\n        self.ln_ff_conv = nn.LayerNorm(encoder_dim)\n\n    def forward(self, x, mask):\n        mask_pad = (mask).long().bool().unsqueeze(1)\n        mask_pad = ~(mask_pad.permute(0, 2, 1) * mask_pad)\n        mask_flat = mask.view(-1).bool()\n        bs, slen, nfeats = x.shape\n\n        # MHSA\n        residual = x\n        x = x * self.scale_mhsa + self.bias_mhsa\n        x = residual + self.mhsa(x, mask_pad)\n\n        # Skip padding for layer norm\n        x_skip = x.view(-1, x.shape[-1])\n        x = x_skip[mask_flat].unsqueeze(0)\n        x = self.ln_mhsa(x)\n\n        # FFN after MHSA\n        residual = x\n        x = x * self.scale_ff_mhsa + self.bias_ff_mhsa\n        x = residual + self.ff_mhsa(x)\n        x = self.ln_ff_mhsa(x)\n\n        # Restore shape\n        x_skip[mask_flat] = x[0].to(dtype=x_skip.dtype)\n        x = x_skip.view(bs, slen, nfeats)\n\n        # Conv\n        residual = x\n        x = x * self.scale_conv + self.bias_conv\n        x = residual + self.conv(x, mask_pad=mask.bool().unsqueeze(1))\n\n        # Skip padding for layer norm\n        x_skip = x.view(-1, x.shape[-1])\n        x = x_skip[mask_flat].unsqueeze(0)\n        x = self.ln_conv(x)\n\n        # FFN after Conv\n        residual = x\n        x = x * self.scale_ff_conv + self.bias_ff_conv\n        x = residual + self.ff_conv(x)\n        x = self.ln_ff_conv(x)\n\n        # Restore shape\n        x_skip[mask_flat] = x[0].to(dtype=x_skip.dtype)\n        x = x_skip.view(bs, slen, nfeats)\n\n        return x\n\nclass SqueezeformerEncoder(nn.Module):\n    \"\"\"\n    Stack of Squeezeformer blocks\n    \"\"\"\n    def __init__(\n        self,\n        input_dim: int = 80,\n        encoder_dim: int = 512,\n        num_layers: int = 16,\n        num_attention_heads: int = 8,\n        feed_forward_expansion_factor: int = 4,\n        conv_expansion_factor: int = 2,\n        input_dropout_p: float = 0.0,\n        feed_forward_dropout_p: float = 0.0,\n        attention_dropout_p: float = 0.0,\n        conv_dropout_p: float = 0.0,\n        conv_kernel_size: int = 31,\n        use_bn: bool = True,\n    ):\n        super(SqueezeformerEncoder, self).__init__()\n        self.num_layers = num_layers\n\n        self.blocks = nn.ModuleList()\n        for idx in range(num_layers):\n            self.blocks.append(\n                SqueezeformerBlock(\n                    encoder_dim=encoder_dim,\n                    num_attention_heads=num_attention_heads,\n                    feed_forward_expansion_factor=feed_forward_expansion_factor,\n                    conv_expansion_factor=conv_expansion_factor,\n                    feed_forward_dropout_p=feed_forward_dropout_p,\n                    attention_dropout_p=attention_dropout_p,\n                    conv_dropout_p=conv_dropout_p,\n                    conv_kernel_size=conv_kernel_size,\n                    use_bn=use_bn,\n                )\n            )\n\n    def forward(self, x: Tensor, mask: Tensor):\n        for idx, block in enumerate(self.blocks):\n            x = block(x, mask)\n        return x\n\nfrom timm.layers.norm_act import BatchNormAct2d\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch_geometric.nn import TransformerConv, LayerNorm\n\nclass SpatialMouseGNN(nn.Module):\n    def __init__(self, in_channels, out_channels, edge_dim=8, num_mice=4, heads=4, dropout=0.1):\n        super().__init__()\n        self.num_mice = num_mice\n        \n        self.hidden_edge_dim = 32  # Give the edges some bandwidth!\n        # 1. Project Raw Node Features\n        self.embedding = nn.Linear(in_channels, out_channels)\n        \n        # 2. Project Raw Edge Features (New)\n        # Projects your 6 social features to match the GNN's internal dimension\n        self.edge_embedding = nn.Sequential(\n            nn.Linear(edge_dim, self.hidden_edge_dim),\n            nn.ReLU(),\n            nn.Linear(self.hidden_edge_dim, self.hidden_edge_dim) # Extra depth\n        )\n        \n        # 3. TransformerConv Layers\n        # Added edge_dim argument so it knows what to expect\n        self.conv1 = TransformerConv(out_channels, out_channels // heads, heads=heads, \n                                     dropout=dropout, beta=True, edge_dim=self.hidden_edge_dim)\n        self.norm1 = LayerNorm(out_channels)\n        \n        self.conv2 = TransformerConv(out_channels, out_channels // heads, heads=heads, \n                                     dropout=dropout, beta=True, edge_dim=self.hidden_edge_dim)\n        self.norm2 = LayerNorm(out_channels)\n        \n        self.dropout = nn.Dropout(dropout)\n\n    def _get_fully_connected_edge_index(self, batch_size_time, device):\n        \"\"\"\n        Creates edges so every mouse connects to every other mouse within the same frame.\n        \"\"\"\n        # Base edges for one frame of 4 mice (0-3)\n        # 0->1, 0->2, 0->3, 1->0, ...\n        base_src = torch.tensor([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], device=device)\n        base_dst = torch.tensor([1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2], device=device)\n        \n        # Offsets for the whole batch\n        offsets = torch.arange(batch_size_time, device=device) * self.num_mice\n        \n        # Broadcasting to create [Batch_Time, Edges_Per_Frame]\n        # We use (BT, 1) + (1, 12) to get (BT, 12)\n        # This ensures edges are ordered: [Graph0_Edges, Graph1_Edges, ...]\n        src = offsets.unsqueeze(1) + base_src.unsqueeze(0)\n        dst = offsets.unsqueeze(1) + base_dst.unsqueeze(0)\n        \n        edge_index = torch.stack([src.flatten(), dst.flatten()], dim=0)\n        return edge_index\n\n    def forward(self, x, x_edges):\n        \"\"\"\n        x: [Batch, Frames, Mice, Node_Features]\n        x_edges: [Batch, Frames, Mice, Mice, Edge_Features]\n        \"\"\"\n        B, T, M, D_n = x.shape\n        _, _, _, _, D_e = x_edges.shape\n        \n        # 1. Flatten Nodes: [Total_Mice, Dim]\n        x_flat = x.view(B * T * M, D_n)\n        x_emb = F.relu(self.embedding(x_flat))\n        \n        # 2. Flatten Edge Attributes to match Edge Index\n        # Reshape to [Batch*Time, M, M, Edge_Dim]\n        edge_feats_flat = x_edges.view(B * T, M, M, D_e)\n        \n        # We need to extract the specific 12 edges (0->1, 0->2...) per graph\n        # Base indices for one graph\n        base_src_idx = torch.tensor([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], device=x.device)\n        base_dst_idx = torch.tensor([1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2], device=x.device)\n        \n        # Select attributes for valid edges\n        # Shape: [Batch*Time, 12, Edge_Dim]\n        valid_edge_attrs = edge_feats_flat[:, base_src_idx, base_dst_idx, :]\n        \n        # Flatten to [Total_Edges, Edge_Dim]\n        # This matches the order of _get_fully_connected_edge_index\n        edge_attr = valid_edge_attrs.view(-1, D_e)\n        \n        # Project Edge Attributes\n        edge_attr = self.edge_embedding(edge_attr)\n\n        # 3. Create Edge Index\n        edge_index = self._get_fully_connected_edge_index(B * T, x.device)\n        \n        # 4. GNN Layers\n        h = self.conv1(x_emb, edge_index, edge_attr=edge_attr)\n        h = self.norm1(h)\n        h = F.relu(h)\n        h = self.dropout(h)\n        \n        h = self.conv2(h, edge_index, edge_attr=edge_attr)\n        h = self.norm2(h)\n        h = F.relu(h)\n        \n        # 5. Reshape back\n        return h.view(B, T, M, -1)\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass RelationAugmentedHead(nn.Module):\n    def __init__(self, encoder_dim, num_classes=39, dropout=0.0):\n        super().__init__()\n        \n        # --- 1. The Relation Network Module (Global Context) ---\n        # g_theta: Processes every pair to find relationships\n        # Input: [Mouse_A + Mouse_B] -> 2 * encoder_dim\n        self.g_theta = nn.Sequential(\n            nn.Linear(encoder_dim * 2, encoder_dim),\n            nn.ReLU(),\n            nn.Linear(encoder_dim, encoder_dim // 2),\n            nn.ReLU()\n        )\n        \n        # f_phi: Processes the SUM of all relations (The \"Cage Vibe\")\n        self.f_phi = nn.Sequential(\n            nn.Linear(encoder_dim // 2, encoder_dim // 2),\n            nn.ReLU(),\n            nn.Dropout(dropout)\n        )\n        \n        # --- 2. The Pairwise Classifier ---\n        # Input: [Mouse_A + Mouse_B + Global_Context]\n        self.in_features = (encoder_dim * 2) + (encoder_dim // 2)\n        \n        self.classifier = nn.Sequential(\n            nn.Linear(self.in_features, encoder_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(encoder_dim, num_classes)\n        )\n\n    def forward(self, x):\n        \"\"\"\n        Args:\n            x: [Batch, Time, Mice, Dim] - Output from Squeezeformer/GNN\n        \"\"\"\n        B, T, M, D = x.shape\n        \n        # --- Step A: Form All Pairs (Same as before) ---\n        # [B, T, M, M, D]\n        m1 = x.unsqueeze(3).expand(-1, -1, -1, M, -1)\n        m2 = x.unsqueeze(2).expand(-1, -1, M, -1, -1)\n        \n        # The specific pairs we want to classify\n        # Shape: [B, T, M, M, 2*D]\n        specific_pairs = torch.cat([m1, m2], dim=-1)\n        \n        # --- Step B: The Relation Network (Global \"Thinking\") ---\n        # 1. Apply g_theta to ALL pairs\n        # Shape: [B, T, M, M, D/2]\n        relations = self.g_theta(specific_pairs)\n        \n        # 2. Sum over all pairs (M*M) to get ONE vector per frame\n        # This represents the \"Global Social State\" of the cage at time T\n        # Shape: [B, T, D/2]\n        global_sum = relations.sum(dim=(2, 3)) \n        \n        # 3. Apply f_phi (Reasoning over the global state)\n        # Shape: [B, T, D/2]\n        global_context = self.f_phi(global_sum)\n        \n        # --- Step C: Combine Local + Global ---\n        # We need to broadcast the global context to every specific pair\n        # Global: [B, T, 1, 1, D/2]\n        global_context_expanded = global_context.unsqueeze(2).unsqueeze(3).expand(-1, -1, M, M, -1)\n        \n        # Concatenate: [Specific_Pair (2D) | Global_Context (0.5D)]\n        # Shape: [B, T, M, M, 2.5D]\n        combined_features = torch.cat([specific_pairs, global_context_expanded], dim=-1)\n        \n        # --- Step D: Final Classification ---\n        # Flatten pairs: [B, T, 16, 2.5D]\n        combined_flat = combined_features.view(B, T, M*M, -1)\n        \n        # Predict\n        logits = self.classifier(combined_flat)\n        \n        return logits\n\nclass Net(nn.Module):\n    \"\"\"\n    Squeezeformer model for MABe mouse behavior detection\n    \"\"\"\n    def __init__(self, cfg):\n        super(Net, self).__init__()\n        self.cfg = cfg\n        # Model dimensions\n        self.encoder_dim = cfg.encoder_config.encoder_dim\n        self.num_pairs = 16  # Number of mouse pairs\n        self.num_actions = 39  # Number of behavior classes (including no_action)\n\n        self.gnn = SpatialMouseGNN(\n            in_channels=cfg.per_mouse_feature_dim,\n            out_channels=cfg.encoder_config.encoder_dim,\n            num_mice=4,\n            heads=4,\n            dropout=cfg.encoder_config.input_dropout_p\n        )\n\n        self.encoder = SqueezeformerEncoder(\n            input_dim=self.encoder_dim,\n            encoder_dim=self.encoder_dim,\n            num_layers=cfg.encoder_config.num_layers,\n            num_attention_heads=cfg.encoder_config.num_attention_heads,\n            feed_forward_expansion_factor=cfg.encoder_config.feed_forward_expansion_factor,\n            conv_expansion_factor=cfg.encoder_config.conv_expansion_factor,\n            input_dropout_p=cfg.encoder_config.input_dropout_p,\n            feed_forward_dropout_p=cfg.encoder_config.feed_forward_dropout_p,\n            attention_dropout_p=cfg.encoder_config.attention_dropout_p,\n            conv_dropout_p=cfg.encoder_config.conv_dropout_p,\n            conv_kernel_size=cfg.encoder_config.conv_kernel_size,\n            use_bn = cfg.use_bn\n        )\n\n        self.classifier = RelationAugmentedHead(\n            encoder_dim=self.encoder_dim,\n            num_classes=self.num_actions,\n            dropout=cfg.encoder_config.feed_forward_dropout_p\n        )\n\n        if hasattr(cfg, 'class_weights') and cfg.class_weights is not None:\n            class_weights = torch.tensor(cfg.class_weights)\n        else:\n            pos_weight = cfg.pos_weight if hasattr(cfg, 'pos_weight') else 50.0 \n            class_weights = torch.full((self.num_actions,), pos_weight) # High weight for all\n            no_action_idx = cfg.action_id_map.get('no_action', -1)\n            if no_action_idx != -1:\n                class_weights[no_action_idx] = 1.0 # Low weight for no_action\n            else:\n                # Fallback if cfg is wrong, assume last class is no_action\n                class_weights[-1] = 1.0\n\n        self.register_buffer('class_weights', class_weights)\n\n        # Use CrossEntropyLoss for multi-class classification\n        self.loss_fn = nn.CrossEntropyLoss(weight=self.class_weights, reduction='none')\n        self.unweighted_ce_fn = nn.CrossEntropyLoss(reduction='none')\n\n        # Optional: Multi-class focal loss\n        self.use_focal_loss = cfg.use_focal_loss if hasattr(cfg, 'use_focal_loss') else False\n        self.focal_gamma = cfg.focal_gamma if hasattr(cfg, 'focal_gamma') else 2.0\n\n        # Training settings\n        self.return_logits = cfg.return_logits if hasattr(cfg, 'return_logits') else False\n\n        print(f'Model initialized with {count_parameters(self):,} trainable parameters')\n        # print(f'Loss: {\"Multi-class Focal\" if self.use_focal_loss else \"Weighted CrossEntropy\"}')\n        # print(f'Class weights: {self.class_weights.cpu().numpy()}')\n\n    def forward(self, batch):\n        mask = batch['input_mask'].long()  # (batch, seq_len)\n        mask_for_encoder = mask.repeat_interleave(4, dim=0)\n\n        if self.cfg.cnn_extractor:\n            x = batch['input_mice']\n            x_nodes = batch['node_feats']  # Shape [B, T, 4, Node_Dim]\n            x_edges = batch['edge_feats']  # Shape [B, T, 4, 4, Edge_Dim]\n            x = self.gnn(x_nodes, x_edges)  # (batch, seq_len, num_mice, encoder_dim)\n            B, T, M, C = x.shape\n            x = x.permute(0,2,1,3).reshape(B * M, T, C)  # (batch*num_mice, seq_len, encoder_dim)\n\n\n        # Encode\n        x = self.encoder(x, mask_for_encoder)  # (batch, seq_len, encoder_dim)\n        x = x.reshape(B, M, T, C).permute(0,2,1,3)  # (batch, seq_len, num_mice, encoder_dim)\n\n        # Classify\n        logits = self.classifier(x)  # (batch, seq_len, num_pairs, num_actions)\n\n        # APPLY BEHAVIOR MASK if provided (for evaluation)\n        if 'behavior_mask' in batch and batch['behavior_mask'] is not None:\n            behavior_mask = batch['behavior_mask']  # (batch, num_pairs, num_actions)\n            # Expand to match logits shape\n            behavior_mask_expanded = behavior_mask.unsqueeze(1).expand_as(logits)\n            # For multi-class, we need to ensure at least one action is valid per pair\n            # Set masked positions to very negative value\n            logits = torch.where(behavior_mask_expanded.bool(), logits,\n                                torch.tensor(-1e10, dtype=logits.dtype, device=logits.device) if logits.dtype == torch.float32\n                                else torch.tensor(-1e4, dtype=logits.dtype, device=logits.device))\n\n        output = {}\n\n        # Calculate loss if labels provided\n        if 'labels' in batch and batch['labels'] is not None:\n            labels = batch['labels']\n\n            # Convert one-hot to class indices if needed\n            if labels.dim() == 4:  # One-hot encoded (batch, seq_len, num_pairs, num_actions)\n                # Convert to class indices\n                labels = torch.argmax(labels, dim=-1)  # (batch, seq_len, num_pairs)\n\n            # Reshape for loss calculation\n            batch_size, seq_len, num_pairs, num_actions = logits.shape\n            logits_flat = logits.reshape(-1, num_actions)  # (batch*seq*pairs, num_actions)\n            labels_flat = labels.reshape(-1)  # (batch*seq*pairs,)\n\n            # Calculate loss\n            if self.use_focal_loss:\n                # Multi-class focal loss\n                unweighted_ce = self.unweighted_ce_fn(logits_flat, labels_flat)\n                pt = torch.exp(-unweighted_ce)\n                alpha = self.class_weights[labels_flat]  # Gather alpha for true classes\n                focal_loss = alpha * (1 - pt) ** self.focal_gamma * unweighted_ce\n                loss = focal_loss\n            else:\n                # Standard weighted cross-entropy\n                loss = self.loss_fn(logits_flat, labels_flat)\n\n            # Reshape loss back\n            loss = loss.view(batch_size, seq_len, num_pairs)\n\n            mask_expanded = mask.unsqueeze(-1).expand_as(loss)\n            loss = loss * mask_expanded\n\n            # Average over valid positions\n            valid_positions = mask_expanded.sum()\n            if valid_positions > 0:\n                loss = loss.sum() / valid_positions\n            else:\n                loss = loss.sum()\n            output['loss'] = loss\n\n        probs = torch.softmax(logits, dim=-1)  # Multi-class probabilities\n        output['predictions'] = probs\n        output['logits'] = logits\n\n        # Store additional info for evaluation\n        if 'video_id' in batch:\n            output['video_id'] = batch['video_id']\n        if 'start_frame' in batch:\n            output['start_frame'] = batch['start_frame']\n\n        return output","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-12-18T13:06:49.064781Z","iopub.execute_input":"2025-12-18T13:06:49.065128Z","iopub.status.idle":"2025-12-18T13:06:49.092555Z","shell.execute_reply.started":"2025-12-18T13:06:49.065103Z","shell.execute_reply":"2025-12-18T13:06:49.091980Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### MODEL 244","metadata":{}},{"cell_type":"code","source":"%%writefile model_244.py\n\n\nimport torch\nfrom torch.nn import functional as F\nfrom torch import nn\nfrom typing import Tuple, Union, Optional\nfrom torch import Tensor\nimport math\nimport numpy as np\n\n\ndef count_parameters(model):\n    \"\"\"Count the number of trainable parameters in model\"\"\"\n    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n\nclass Swish(nn.Module):\n    \"\"\"Swish activation function\"\"\"\n    def __init__(self) -> None:\n        super(Swish, self).__init__()\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return inputs * inputs.sigmoid()\n\nclass GLU(nn.Module):\n    \"\"\"Gated Linear Unit activation\"\"\"\n    def __init__(self, dim: int) -> None:\n        super(GLU, self).__init__()\n        self.dim = dim\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        outputs, gate = inputs.chunk(2, dim=self.dim)\n        return outputs * gate.sigmoid()\n\nclass FeedForwardModule(nn.Module):\n    \"\"\"\n    Feed Forward Module with pre-norm residual units\n    \"\"\"\n    def __init__(\n        self,\n        encoder_dim: int = 512,\n        expansion_factor: int = 4,\n        dropout_p: float = 0.0,\n    ) -> None:\n        super(FeedForwardModule, self).__init__()\n\n        self.ffn1 = nn.Linear(encoder_dim, encoder_dim * expansion_factor, bias=True)\n        self.act = Swish()\n        self.do1 = nn.Dropout(p=dropout_p)\n        self.ffn2 = nn.Linear(encoder_dim * expansion_factor, encoder_dim, bias=True)\n        self.do2 = nn.Dropout(p=dropout_p)\n\n    def forward(self, x):\n        x = self.ffn1(x)\n        x = self.act(x)\n        x = self.do1(x)\n        x = self.ffn2(x)\n        x = self.do2(x)\n        return x\n\nclass RelPositionalEncoding(nn.Module):\n    \"\"\"\n    Relative positional encoding module for handling variable sequence lengths\n    \"\"\"\n    def __init__(self, d_model: int = 512, max_len: int = 5000) -> None:\n        super(RelPositionalEncoding, self).__init__()\n        self.d_model = d_model\n        self.pe = None\n        self.extend_pe(torch.tensor(0.0).expand(1, max_len))\n\n    def extend_pe(self, x):\n        if self.pe is not None:\n            if self.pe.size(1) >= x.size(1) * 2 - 1:\n                if self.pe.dtype != x.dtype or self.pe.device != x.device:\n                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)\n                return\n\n        pe_positive = torch.zeros(x.size(1), self.d_model)\n        pe_negative = torch.zeros(x.size(1), self.d_model)\n        position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)\n        div_term = torch.exp(\n            torch.arange(0, self.d_model, 2, dtype=torch.float32) * -(math.log(10000.0) / self.d_model)\n        )\n        pe_positive[:, 0::2] = torch.sin(position * div_term)\n        pe_positive[:, 1::2] = torch.cos(position * div_term)\n        pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)\n        pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)\n\n        pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)\n        pe_negative = pe_negative[1:].unsqueeze(0)\n        pe = torch.cat([pe_positive, pe_negative], dim=1)\n        self.pe = pe.to(device=x.device, dtype=x.dtype)\n\n    def forward(self, x: torch.Tensor):\n        self.extend_pe(x)\n        pos_emb = self.pe[\n            :,\n            self.pe.size(1) // 2 - x.size(1) + 1 : self.pe.size(1) // 2 + x.size(1),\n        ]\n        return pos_emb\n\nclass RelativeMultiHeadAttention(nn.Module):\n    \"\"\"\n    Multi-head attention with relative positional encoding\n    \"\"\"\n    def __init__(\n        self,\n        d_model: int = 512,\n        num_heads: int = 16,\n        dropout_p: float = 0.0,\n    ):\n        super(RelativeMultiHeadAttention, self).__init__()\n        assert d_model % num_heads == 0, \"d_model % num_heads should be zero.\"\n        self.d_model = d_model\n        self.d_head = int(d_model / num_heads)\n        self.num_heads = num_heads\n        self.sqrt_dim = math.sqrt(self.d_head)\n\n        self.query_proj = nn.Linear(d_model, d_model)\n        self.key_proj = nn.Linear(d_model, d_model)\n        self.value_proj = nn.Linear(d_model, d_model)\n        self.pos_proj = nn.Linear(d_model, d_model, bias=False)\n\n        self.dropout = nn.Dropout(p=dropout_p)\n        self.u_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))\n        self.v_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))\n        torch.nn.init.xavier_uniform_(self.u_bias)\n        torch.nn.init.xavier_uniform_(self.v_bias)\n\n        self.out_proj = nn.Linear(d_model, d_model)\n\n    def forward(\n        self,\n        query: Tensor,\n        key: Tensor,\n        value: Tensor,\n        pos_embedding: Tensor,\n        mask: Optional[Tensor] = None,\n    ) -> Tensor:\n        batch_size = value.size(0)\n\n        query = self.query_proj(query).view(batch_size, -1, self.num_heads, self.d_head)\n        key = self.key_proj(key).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)\n        value = self.value_proj(value).view(batch_size, -1, self.num_heads, self.d_head).permute(0, 2, 1, 3)\n        pos_embedding = self.pos_proj(pos_embedding).view(batch_size, -1, self.num_heads, self.d_head)\n\n        content_score = torch.matmul((query + self.u_bias).transpose(1, 2), key.transpose(2, 3))\n        pos_score = torch.matmul((query + self.v_bias).transpose(1, 2), pos_embedding.permute(0, 2, 3, 1))\n        pos_score = self._relative_shift(pos_score)\n\n        score = (content_score + pos_score) / self.sqrt_dim\n\n        if mask is not None:\n            mask = mask.unsqueeze(1)\n            score.masked_fill_(mask, -1e9) if score.dtype == torch.float32 else score.masked_fill_(mask, -1e4)\n\n        attn = F.softmax(score, -1)\n        attn = self.dropout(attn)\n\n        context = torch.matmul(attn, value).transpose(1, 2)\n        context = context.contiguous().view(batch_size, -1, self.d_model)\n\n        return self.out_proj(context)\n\n    def _relative_shift(self, pos_score: Tensor) -> Tensor:\n        batch_size, num_heads, seq_length1, seq_length2 = pos_score.size()\n        zeros = pos_score.new_zeros(batch_size, num_heads, seq_length1, 1)\n        padded_pos_score = torch.cat([zeros, pos_score], dim=-1)\n\n        padded_pos_score = padded_pos_score.view(batch_size, num_heads, seq_length2 + 1, seq_length1)\n        pos_score = padded_pos_score[:, :, 1:].view_as(pos_score)[:, :, :, : seq_length2 // 2 + 1]\n\n        return pos_score\n\nclass MultiHeadedSelfAttentionModule(nn.Module):\n    \"\"\"\n    Self-attention module with relative positional encoding\n    \"\"\"\n    def __init__(self, d_model: int, num_heads: int, dropout_p: float = 0.0):\n        super(MultiHeadedSelfAttentionModule, self).__init__()\n        self.positional_encoding = RelPositionalEncoding(d_model)\n        self.attention = RelativeMultiHeadAttention(d_model, num_heads, dropout_p)\n        self.dropout = nn.Dropout(p=dropout_p)\n\n    def forward(self, inputs: Tensor, mask: Optional[Tensor] = None):\n        batch_size = inputs.size(0)\n        pos_embedding = self.positional_encoding(inputs)\n        pos_embedding = pos_embedding.repeat(batch_size, 1, 1)\n\n        outputs = self.attention(inputs, inputs, inputs, pos_embedding=pos_embedding, mask=mask)\n        return self.dropout(outputs)\n\nclass DepthwiseConv1d(nn.Module):\n    \"\"\"Depthwise 1D convolution\"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        kernel_size: int,\n        stride: int = 1,\n        padding: int = 0,\n        bias: bool = False,\n    ) -> None:\n        super(DepthwiseConv1d, self).__init__()\n        assert out_channels % in_channels == 0, \"out_channels should be constant multiple of in_channels\"\n        self.conv = nn.Conv1d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=kernel_size,\n            groups=in_channels,\n            stride=stride,\n            padding=padding,\n            bias=bias,\n        )\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return self.conv(inputs)\n\nclass PointwiseConv1d(nn.Module):\n    \"\"\"Pointwise 1D convolution (kernel size = 1)\"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        out_channels: int,\n        stride: int = 1,\n        padding: int = 0,\n        bias: bool = True,\n    ) -> None:\n        super(PointwiseConv1d, self).__init__()\n        self.conv = nn.Conv1d(\n            in_channels=in_channels,\n            out_channels=out_channels,\n            kernel_size=1,\n            stride=stride,\n            padding=padding,\n            bias=bias,\n        )\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        return self.conv(inputs)\n\nclass ConvModule(nn.Module):\n    \"\"\"\n    Convolution module with pointwise conv -> GLU -> depthwise conv -> normalization -> activation\n    \"\"\"\n    def __init__(\n        self,\n        in_channels: int,\n        kernel_size: int = 31,\n        expansion_factor: int = 2,\n        dropout_p: float = 0.0,\n        use_bn: bool = True,\n    ) -> None:\n        super(ConvModule, self).__init__()\n        assert (kernel_size - 1) % 2 == 0, \"kernel_size should be a odd number for 'SAME' padding\"\n        assert expansion_factor == 2, \"Currently, Only Supports expansion_factor 2\"\n\n        self.pw_conv_1 = PointwiseConv1d(in_channels, in_channels * expansion_factor, stride=1, padding=0, bias=True)\n        self.act1 = GLU(dim=1)\n        self.dw_conv = DepthwiseConv1d(in_channels, in_channels, kernel_size, stride=1, padding=(kernel_size - 1) // 2)\n        self.bn = nn.BatchNorm1d(in_channels)\n        self.inorm = nn.InstanceNorm1d(in_channels, affine=True)\n        self.act2 = Swish()\n        self.pw_conv_2 = PointwiseConv1d(in_channels, in_channels, stride=1, padding=0, bias=True)\n        self.do = nn.Dropout(p=dropout_p)\n        self.use_bn = use_bn\n\n    def forward(self, x, mask_pad):\n        # Transpose for conv operations [B, T, C]\n        x = x.transpose(1, 2)\n        if mask_pad.size(2) > 0:  # time > 0\n            x = x.masked_fill(~mask_pad, 0.0)\n\n        x = self.pw_conv_1(x)\n        x = self.act1(x)\n        x = self.dw_conv(x)\n\n        if self.use_bn:\n            # Apply batch norm only to non-padded positions\n            x_bn = x.permute(0,2,1).reshape(-1, x.shape[1])\n            mask_bn = mask_pad.view(-1)\n            x_bn[mask_bn] = self.bn(x_bn[mask_bn])\n            x = x_bn.view(x.permute(0,2,1).shape).permute(0,2,1)\n        else:    \n            x = self.inorm(x)\n\n        x = self.act2(x)\n        x = self.pw_conv_2(x)\n        x = self.do(x)\n\n        # Mask batch padding again\n        if mask_pad.size(2) > 0:  # time > 0\n            x = x.masked_fill(~mask_pad, 0.0)\n        x = x.transpose(1, 2)\n        return x\n\ndef make_scale(encoder_dim):\n    \"\"\"Create learnable scale and bias parameters\"\"\"\n    scale = torch.nn.Parameter(torch.tensor([1.] * encoder_dim)[None, None, :])\n    bias = torch.nn.Parameter(torch.tensor([0.] * encoder_dim)[None, None, :])\n    return scale, bias\n\nclass SqueezeformerBlock(nn.Module):\n    \"\"\"\n    Squeezeformer block: MHSA -> FFN -> Conv -> FFN with residual connections\n    \"\"\"\n    def __init__(\n        self,\n        encoder_dim: int = 512,\n        num_attention_heads: int = 8,\n        feed_forward_expansion_factor: int = 4,\n        conv_expansion_factor: int = 2,\n        feed_forward_dropout_p: float = 0.0,\n        attention_dropout_p: float = 0.1,\n        conv_dropout_p: float = 0.1,\n        conv_kernel_size: int = 31,\n        use_bn: bool = True,\n    ):\n        super(SqueezeformerBlock, self).__init__()\n\n        self.scale_mhsa, self.bias_mhsa = make_scale(encoder_dim)\n        self.scale_ff_mhsa, self.bias_ff_mhsa = make_scale(encoder_dim)\n        self.scale_conv, self.bias_conv = make_scale(encoder_dim)\n        self.scale_ff_conv, self.bias_ff_conv = make_scale(encoder_dim)\n\n        self.mhsa = MultiHeadedSelfAttentionModule(\n            d_model=encoder_dim,\n            num_heads=num_attention_heads,\n            dropout_p=attention_dropout_p,\n        )\n        self.ln_mhsa = nn.LayerNorm(encoder_dim)\n        self.ff_mhsa = FeedForwardModule(\n            encoder_dim=encoder_dim,\n            expansion_factor=feed_forward_expansion_factor,\n            dropout_p=feed_forward_dropout_p,\n        )\n        self.ln_ff_mhsa = nn.LayerNorm(encoder_dim)\n        self.conv = ConvModule(\n            in_channels=encoder_dim,\n            kernel_size=conv_kernel_size,\n            expansion_factor=conv_expansion_factor,\n            dropout_p=conv_dropout_p,\n            use_bn=use_bn,\n        )\n        self.ln_conv = nn.LayerNorm(encoder_dim)\n        self.ff_conv = FeedForwardModule(\n            encoder_dim=encoder_dim,\n            expansion_factor=feed_forward_expansion_factor,\n            dropout_p=feed_forward_dropout_p,\n        )\n        self.ln_ff_conv = nn.LayerNorm(encoder_dim)\n\n    def forward(self, x, mask):\n        mask_pad = (mask).long().bool().unsqueeze(1)\n        mask_pad = ~(mask_pad.permute(0, 2, 1) * mask_pad)\n        mask_flat = mask.view(-1).bool()\n        bs, slen, nfeats = x.shape\n\n        # MHSA\n        residual = x\n        x = x * self.scale_mhsa + self.bias_mhsa\n        x = residual + self.mhsa(x, mask_pad)\n\n        # Skip padding for layer norm\n        x_skip = x.view(-1, x.shape[-1])\n        x = x_skip[mask_flat].unsqueeze(0)\n        x = self.ln_mhsa(x)\n\n        # FFN after MHSA\n        residual = x\n        x = x * self.scale_ff_mhsa + self.bias_ff_mhsa\n        x = residual + self.ff_mhsa(x)\n        x = self.ln_ff_mhsa(x)\n\n        # Restore shape\n        x_skip[mask_flat] = x[0].to(dtype=x_skip.dtype)\n        x = x_skip.view(bs, slen, nfeats)\n\n        # Conv\n        residual = x\n        x = x * self.scale_conv + self.bias_conv\n        x = residual + self.conv(x, mask_pad=mask.bool().unsqueeze(1))\n\n        # Skip padding for layer norm\n        x_skip = x.view(-1, x.shape[-1])\n        x = x_skip[mask_flat].unsqueeze(0)\n        x = self.ln_conv(x)\n\n        # FFN after Conv\n        residual = x\n        x = x * self.scale_ff_conv + self.bias_ff_conv\n        x = residual + self.ff_conv(x)\n        x = self.ln_ff_conv(x)\n\n        # Restore shape\n        x_skip[mask_flat] = x[0].to(dtype=x_skip.dtype)\n        x = x_skip.view(bs, slen, nfeats)\n\n        return x\n\nclass SqueezeformerEncoder(nn.Module):\n    \"\"\"\n    Stack of Squeezeformer blocks\n    \"\"\"\n    def __init__(\n        self,\n        input_dim: int = 80,\n        encoder_dim: int = 512,\n        num_layers: int = 16,\n        num_attention_heads: int = 8,\n        feed_forward_expansion_factor: int = 4,\n        conv_expansion_factor: int = 2,\n        input_dropout_p: float = 0.0,\n        feed_forward_dropout_p: float = 0.0,\n        attention_dropout_p: float = 0.0,\n        conv_dropout_p: float = 0.0,\n        conv_kernel_size: int = 31,\n        use_bn: bool = True,\n    ):\n        super(SqueezeformerEncoder, self).__init__()\n        self.num_layers = num_layers\n\n        self.blocks = nn.ModuleList()\n        for idx in range(num_layers):\n            self.blocks.append(\n                SqueezeformerBlock(\n                    encoder_dim=encoder_dim,\n                    num_attention_heads=num_attention_heads,\n                    feed_forward_expansion_factor=feed_forward_expansion_factor,\n                    conv_expansion_factor=conv_expansion_factor,\n                    feed_forward_dropout_p=feed_forward_dropout_p,\n                    attention_dropout_p=attention_dropout_p,\n                    conv_dropout_p=conv_dropout_p,\n                    conv_kernel_size=conv_kernel_size,\n                    use_bn=use_bn,\n                )\n            )\n\n    def forward(self, x: Tensor, mask: Tensor):\n        for idx, block in enumerate(self.blocks):\n            x = block(x, mask)\n        return x\n\nfrom timm.layers.norm_act import BatchNormAct2d\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch_geometric.nn import TransformerConv, LayerNorm\n\nclass SpatialMouseGNN(nn.Module):\n    def __init__(self, in_channels, out_channels, num_mice=4, heads=4, dropout=0.1):\n        super().__init__()\n        self.num_mice = num_mice\n        # 1. Project Raw Features (Crucial Step)\n        # We need to turn coordinates into \"semantic\" embeddings before Attention\n        self.embedding = nn.Linear(in_channels, out_channels)\n        \n        # 2. TransformerConv Layers\n        # We use beta=True for a gating mechanism (improves deep GNNs)\n        self.conv1 = TransformerConv(out_channels, out_channels // heads, heads=heads, \n                                     dropout=dropout, beta=True)\n        self.norm1 = LayerNorm(out_channels) # LayerNorm is often more stable for Transformers\n        \n        self.conv2 = TransformerConv(out_channels, out_channels // heads, heads=heads, \n                                     dropout=dropout, beta=True)\n        self.norm2 = LayerNorm(out_channels)\n        \n        self.dropout = nn.Dropout(dropout)\n\n    def _get_fully_connected_edge_index(self, batch_size_time, device):\n        \"\"\"\n        Creates edges so every mouse connects to every other mouse within the same frame.\n        \"\"\"\n        # Base edges for one frame of 4 mice (0-3)\n        # 0->1, 0->2, 0->3, 1->0, ...\n        base_src = torch.tensor([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], device=device)\n        base_dst = torch.tensor([1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2], device=device)\n        # If you want self-loops (mice look at themselves), add them here, \n        # but TransformerConv usually adds them automatically or handles them separately.\n\n        # Repeat for the whole batch\n        # We need to offset indices by num_mice (4) for each subsequent frame\n        offsets = torch.arange(batch_size_time, device=device) * self.num_mice\n        \n        # Broadcasting to create the massive edge list\n        # shape: [Edges_Per_Frame, Batch_Time]\n        src = base_src.unsqueeze(1) + offsets.unsqueeze(0)\n        dst = base_dst.unsqueeze(1) + offsets.unsqueeze(0)\n        \n        edge_index = torch.stack([src.flatten(), dst.flatten()], dim=0)\n        return edge_index\n\n    def forward(self, x):\n        \"\"\"\n        Input: [Batch, Frames, Mice, Features]\n        Output: [Batch, Frames, Mice, Out_Channels]\n        \"\"\"\n        B, T, M, D = x.shape\n        \n        # Flatten Batch and Time: treat every frame as an independent graph\n        # New shape: [Total_Graphs * Mice, Features]\n        x_flat = x.view(B * T * M, D)\n        \n        # 1. Linear Projection\n        x_emb = F.relu(self.embedding(x_flat))\n        \n        # 2. Create Edges on the fly\n        # (Optimized: we only calculate this once per forward pass structure)\n        edge_index = self._get_fully_connected_edge_index(B * T, x.device)\n        \n        # 3. GNN Layer 1\n        # TransformerConv expects [Num_Nodes, Dim]\n        h = self.conv1(x_emb, edge_index)\n        h = self.norm1(h)\n        h = F.relu(h)\n        h = self.dropout(h)\n        \n        # 4. GNN Layer 2\n        h = self.conv2(h, edge_index)\n        h = self.norm2(h)\n        h = F.relu(h)\n        \n        # 5. Reshape back (NO POOLING)\n        # We want to keep the mice separate for the Squeezeformer\n        return h.view(B, T, M, -1)\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass RelationAugmentedHead(nn.Module):\n    def __init__(self, encoder_dim, num_classes=39, dropout=0.0):\n        super().__init__()\n        \n        # --- 1. The Relation Network Module (Global Context) ---\n        # g_theta: Processes every pair to find relationships\n        # Input: [Mouse_A + Mouse_B] -> 2 * encoder_dim\n        self.g_theta = nn.Sequential(\n            nn.Linear(encoder_dim * 2, encoder_dim),\n            nn.ReLU(),\n            nn.Linear(encoder_dim, encoder_dim // 2),\n            nn.ReLU()\n        )\n        \n        # f_phi: Processes the SUM of all relations (The \"Cage Vibe\")\n        self.f_phi = nn.Sequential(\n            nn.Linear(encoder_dim // 2, encoder_dim // 2),\n            nn.ReLU(),\n            nn.Dropout(dropout)\n        )\n        \n        # --- 2. The Pairwise Classifier ---\n        # Input: [Mouse_A + Mouse_B + Global_Context]\n        self.in_features = (encoder_dim * 2) + (encoder_dim // 2)\n        \n        self.classifier = nn.Sequential(\n            nn.Linear(self.in_features, encoder_dim),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(encoder_dim, num_classes)\n        )\n\n    def forward(self, x):\n        \"\"\"\n        Args:\n            x: [Batch, Time, Mice, Dim] - Output from Squeezeformer/GNN\n        \"\"\"\n        B, T, M, D = x.shape\n        \n        # --- Step A: Form All Pairs (Same as before) ---\n        # [B, T, M, M, D]\n        m1 = x.unsqueeze(3).expand(-1, -1, -1, M, -1)\n        m2 = x.unsqueeze(2).expand(-1, -1, M, -1, -1)\n        \n        # The specific pairs we want to classify\n        # Shape: [B, T, M, M, 2*D]\n        specific_pairs = torch.cat([m1, m2], dim=-1)\n        \n        # --- Step B: The Relation Network (Global \"Thinking\") ---\n        # 1. Apply g_theta to ALL pairs\n        # Shape: [B, T, M, M, D/2]\n        relations = self.g_theta(specific_pairs)\n        \n        # 2. Sum over all pairs (M*M) to get ONE vector per frame\n        # This represents the \"Global Social State\" of the cage at time T\n        # Shape: [B, T, D/2]\n        global_sum = relations.sum(dim=(2, 3)) \n        \n        # 3. Apply f_phi (Reasoning over the global state)\n        # Shape: [B, T, D/2]\n        global_context = self.f_phi(global_sum)\n        \n        # --- Step C: Combine Local + Global ---\n        # We need to broadcast the global context to every specific pair\n        # Global: [B, T, 1, 1, D/2]\n        global_context_expanded = global_context.unsqueeze(2).unsqueeze(3).expand(-1, -1, M, M, -1)\n        \n        # Concatenate: [Specific_Pair (2D) | Global_Context (0.5D)]\n        # Shape: [B, T, M, M, 2.5D]\n        combined_features = torch.cat([specific_pairs, global_context_expanded], dim=-1)\n        \n        # --- Step D: Final Classification ---\n        # Flatten pairs: [B, T, 16, 2.5D]\n        combined_flat = combined_features.view(B, T, M*M, -1)\n        \n        # Predict\n        logits = self.classifier(combined_flat)\n        \n        return logits\n\nclass Net(nn.Module):\n    \"\"\"\n    Squeezeformer model for MABe mouse behavior detection\n    \"\"\"\n    def __init__(self, cfg):\n        super(Net, self).__init__()\n        self.cfg = cfg\n        # Model dimensions\n        self.encoder_dim = cfg.encoder_config.encoder_dim\n        self.num_pairs = 16  # Number of mouse pairs\n        self.num_actions = 39  # Number of behavior classes (including no_action)\n\n        self.gnn = SpatialMouseGNN(\n            in_channels=cfg.per_mouse_feature_dim,\n            out_channels=cfg.encoder_config.encoder_dim,\n            num_mice=4,\n            heads=4,\n            dropout=cfg.encoder_config.input_dropout_p\n        )\n\n        self.encoder = SqueezeformerEncoder(\n            input_dim=self.encoder_dim,\n            encoder_dim=self.encoder_dim,\n            num_layers=cfg.encoder_config.num_layers,\n            num_attention_heads=cfg.encoder_config.num_attention_heads,\n            feed_forward_expansion_factor=cfg.encoder_config.feed_forward_expansion_factor,\n            conv_expansion_factor=cfg.encoder_config.conv_expansion_factor,\n            input_dropout_p=cfg.encoder_config.input_dropout_p,\n            feed_forward_dropout_p=cfg.encoder_config.feed_forward_dropout_p,\n            attention_dropout_p=cfg.encoder_config.attention_dropout_p,\n            conv_dropout_p=cfg.encoder_config.conv_dropout_p,\n            conv_kernel_size=cfg.encoder_config.conv_kernel_size,\n            use_bn = cfg.use_bn\n        )\n\n        self.classifier = RelationAugmentedHead(\n            encoder_dim=self.encoder_dim,\n            num_classes=self.num_actions,\n            dropout=cfg.encoder_config.feed_forward_dropout_p\n        )\n\n        if hasattr(cfg, 'class_weights') and cfg.class_weights is not None:\n            class_weights = torch.tensor(cfg.class_weights)\n        else:\n            pos_weight = cfg.pos_weight if hasattr(cfg, 'pos_weight') else 50.0 \n            class_weights = torch.full((self.num_actions,), pos_weight) # High weight for all\n            no_action_idx = cfg.action_id_map.get('no_action', -1)\n            if no_action_idx != -1:\n                class_weights[no_action_idx] = 1.0 # Low weight for no_action\n            else:\n                # Fallback if cfg is wrong, assume last class is no_action\n                class_weights[-1] = 1.0\n\n        self.register_buffer('class_weights', class_weights)\n\n        # Use CrossEntropyLoss for multi-class classification\n        self.loss_fn = nn.CrossEntropyLoss(weight=self.class_weights, reduction='none')\n        self.unweighted_ce_fn = nn.CrossEntropyLoss(reduction='none')\n\n        # Optional: Multi-class focal loss\n        self.use_focal_loss = cfg.use_focal_loss if hasattr(cfg, 'use_focal_loss') else False\n        self.focal_gamma = cfg.focal_gamma if hasattr(cfg, 'focal_gamma') else 2.0\n\n        # Training settings\n        self.return_logits = cfg.return_logits if hasattr(cfg, 'return_logits') else False\n\n        print(f'Model initialized with {count_parameters(self):,} trainable parameters')\n        # print(f'Loss: {\"Multi-class Focal\" if self.use_focal_loss else \"Weighted CrossEntropy\"}')\n        # print(f'Class weights: {self.class_weights.cpu().numpy()}')\n\n    def forward(self, batch):\n        mask = batch['input_mask'].long()  # (batch, seq_len)\n        mask_for_encoder = mask.repeat_interleave(4, dim=0)\n\n        if self.cfg.cnn_extractor:\n            x = batch['input_ego_mice']\n            x = self.gnn(x)  # (batch, seq_len, num_mice, encoder_dim)\n            B, T, M, C = x.shape\n            x = x.permute(0,2,1,3).reshape(B * M, T, C)  # (batch*num_mice, seq_len, encoder_dim)\n\n\n        # Encode\n        x = self.encoder(x, mask_for_encoder)  # (batch, seq_len, encoder_dim)\n        x = x.reshape(B, M, T, C).permute(0,2,1,3)  # (batch, seq_len, num_mice, encoder_dim)\n\n        # Classify\n        logits = self.classifier(x)  # (batch, seq_len, num_pairs, num_actions)\n\n        # APPLY BEHAVIOR MASK if provided (for evaluation)\n        if 'behavior_mask' in batch and batch['behavior_mask'] is not None:\n            behavior_mask = batch['behavior_mask']  # (batch, num_pairs, num_actions)\n            # Expand to match logits shape\n            behavior_mask_expanded = behavior_mask.unsqueeze(1).expand_as(logits)\n            # For multi-class, we need to ensure at least one action is valid per pair\n            # Set masked positions to very negative value\n            logits = torch.where(behavior_mask_expanded.bool(), logits,\n                                torch.tensor(-1e10, dtype=logits.dtype, device=logits.device) if logits.dtype == torch.float32\n                                else torch.tensor(-1e4, dtype=logits.dtype, device=logits.device))\n\n        output = {}\n\n        # Calculate loss if labels provided\n        if 'labels' in batch and batch['labels'] is not None:\n            labels = batch['labels']\n\n            # Convert one-hot to class indices if needed\n            if labels.dim() == 4:  # One-hot encoded (batch, seq_len, num_pairs, num_actions)\n                # Convert to class indices\n                labels = torch.argmax(labels, dim=-1)  # (batch, seq_len, num_pairs)\n\n            # Reshape for loss calculation\n            batch_size, seq_len, num_pairs, num_actions = logits.shape\n            logits_flat = logits.reshape(-1, num_actions)  # (batch*seq*pairs, num_actions)\n            labels_flat = labels.reshape(-1)  # (batch*seq*pairs,)\n\n            # Calculate loss\n            if self.use_focal_loss:\n                # Multi-class focal loss\n                unweighted_ce = self.unweighted_ce_fn(logits_flat, labels_flat)\n                pt = torch.exp(-unweighted_ce)\n                alpha = self.class_weights[labels_flat]  # Gather alpha for true classes\n                focal_loss = alpha * (1 - pt) ** self.focal_gamma * unweighted_ce\n                loss = focal_loss\n            else:\n                # Standard weighted cross-entropy\n                loss = self.loss_fn(logits_flat, labels_flat)\n\n            # Reshape loss back\n            loss = loss.view(batch_size, seq_len, num_pairs)\n\n            mask_expanded = mask.unsqueeze(-1).expand_as(loss)\n            loss = loss * mask_expanded\n\n            # Average over valid positions\n            valid_positions = mask_expanded.sum()\n            if valid_positions > 0:\n                loss = loss.sum() / valid_positions\n            else:\n                loss = loss.sum()\n            output['loss'] = loss\n\n        probs = torch.softmax(logits, dim=-1)  # Multi-class probabilities\n        output['predictions'] = probs\n        output['logits'] = logits\n\n        # Store additional info for evaluation\n        if 'video_id' in batch:\n            output['video_id'] = batch['video_id']\n        if 'start_frame' in batch:\n            output['start_frame'] = batch['start_frame']\n\n        return output","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2025-12-18T13:06:49.093544Z","iopub.execute_input":"2025-12-18T13:06:49.093806Z","iopub.status.idle":"2025-12-18T13:06:49.122295Z","shell.execute_reply.started":"2025-12-18T13:06:49.093781Z","shell.execute_reply":"2025-12-18T13:06:49.121517Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### DATASET","metadata":{}},{"cell_type":"code","source":"import os\nimport torch\nimport numpy as np\nfrom torch.utils.data import Dataset\nimport pandas as pd\nimport torch.nn.functional as F\nimport random\nfrom tqdm import tqdm\nimport ast\nimport warnings\nimport gc\n\nwarnings.filterwarnings('ignore', message='Mean of empty slice')\n\n\ndef batch_to_device(batch, device):\n    batch_dict = {}\n    for key in batch:\n        if torch.is_tensor(batch[key]):\n            batch_dict[key] = batch[key].to(device)\n        else:\n            batch_dict[key] = batch[key]\n    return batch_dict\n\n\nclass CustomDataset(Dataset):\n    def __init__(self, df, cfg, aug=None, mode=\"train\"):\n        self.cfg = cfg\n        self.df = df.copy()\n        self.mode = mode\n        self.aug = aug\n\n        self.base_data_dir = '/kaggle/input/MABe-mouse-behavior-detection'\n        valid_videos = []\n        self.video_info = {}\n        self.video_keypoints = {}\n        \n        num_mice = len(cfg.set_mice)\n        num_pairs = num_mice * num_mice\n        num_actions = len(cfg.set_behavior_classes)\n        mouse_to_idx = {'mouse1': 0, 'mouse2': 1, 'mouse3': 2, 'mouse4': 3}\n\n        for idx in tqdm(range(len(self.df)), desc=\"Scanning videos and pre-loading labels\"):\n            row = self.df.iloc[idx]\n            lab_id = row['lab_id']\n            video_id = row['video_id']\n\n            try:\n                keypoints, num_frames = process_tracking_video(row, 'test')\n                if keypoints is None or num_frames == 0:\n                    continue\n\n                if video_id not in self.video_keypoints.keys():\n                    self.video_keypoints[video_id] = keypoints\n                else:\n                    self.video_keypoints[video_id+len(self.video_keypoints)] = keypoints\n                    print(len(self.video_keypoints), self.video_keypoints.keys())\n                \n                behaviors_labeled = ast.literal_eval(row['behaviors_labeled']) if isinstance(row['behaviors_labeled'], str) else row['behaviors_labeled']\n                if not isinstance(behaviors_labeled, list):\n                    continue\n                behavior_mask = np.zeros((num_pairs, num_actions), dtype=bool)\n                for behavior_str in behaviors_labeled:\n                    parts = behavior_str.split(',')\n                    if len(parts) != 3:\n                        continue\n                    agent, target, action = parts\n                    agent_idx = mouse_to_idx.get(agent, -1)\n                    if target == 'self':\n                        target_idx = agent_idx\n                    else:\n                        target_idx = mouse_to_idx.get(target, -1)\n                    if agent_idx == -1 or target_idx == -1:\n                        continue\n                    pair_idx = agent_idx * num_mice + target_idx\n                    action_idx = cfg.action_id_map.get(action, -1)\n                    if action_idx == -1:\n                        continue\n                    behavior_mask[pair_idx, action_idx] = True\n                \n                no_action_idx = cfg.action_id_map.get('no_action', -1)\n                if no_action_idx != -1:\n                    behavior_mask[:, no_action_idx] = True\n\n                self.video_info[video_id] = {\n                    'num_frames': num_frames,\n                    'duration_sec': row.get('video_duration_sec', num_frames / 30),\n                    'lab_id': row['lab_id'],\n                    'behavior_mask': behavior_mask\n                }\n            except Exception as e:\n                print(f\"Error processing video {video_id}: {e}\")\n                continue\n\n        self.create_sampling_schedule()\n\n    def create_sampling_schedule(self):\n        self.sampling_schedule = []\n\n        for idx in range(len(self.df)):\n            row = self.df.iloc[idx]\n            video_id = row['video_id']\n            lab_id = row['lab_id']\n            if video_id not in self.video_info:\n                continue\n\n            num_frames = self.video_info[video_id]['num_frames']\n            if num_frames <= self.cfg.window_size:\n                num_windows = 1\n            else:\n                stride = self.cfg.stride if hasattr(self.cfg, 'stride') else self.cfg.window_size // 2\n                total_possible_windows = (num_frames // stride ) + 1\n\n                if self.mode == 'train':\n                    num_windows = max(\n                        self.cfg.min_windows_per_video,\n                        min(\n                            self.cfg.max_windows_per_video,\n                            int(total_possible_windows * self.cfg.windows_per_epoch_ratio)\n                        )\n                    )\n                else:\n                    num_windows = total_possible_windows\n\n            # Add entries to schedule\n            for win_idx in range(num_windows):\n                self.sampling_schedule.append({\n                    'video_id': video_id,\n                    'df_idx': idx,\n                    'win_idx': win_idx,\n                    'num_frames': num_frames\n                })\n\n        print(f\"Created sampling schedule with {len(self.sampling_schedule)} windows\")\n        if len(self.df) > 0:\n            print(f\"Average windows per video: {len(self.sampling_schedule) / len(self.df):.2f}\")\n\n    def __len__(self):\n        return len(self.sampling_schedule)\n\n    def __getitem__(self, idx):\n        sample_info = self.sampling_schedule[idx]\n        video_id = sample_info['video_id']\n        num_frames = sample_info['num_frames']\n\n        row = self.df.iloc[sample_info['df_idx']]\n        # Get video info\n        video_data = self.video_info[video_id]\n        \n        if self.mode == 'train':\n            positive_frames = video_data.get('behavior_frames', np.array([], dtype=np.int64))\n            \n            if random.random() < self.cfg.bias_prob and len(positive_frames) > 0:\n                center_frame = random.choice(positive_frames)\n                offset = random.randint(0, self.cfg.window_size - 1)\n                start = max(0, min(center_frame - offset, num_frames - self.cfg.window_size))\n            else:\n                start = random.randint(0, max(0, num_frames - self.cfg.window_size))\n        else:  # val mode: deterministic sliding window\n            stride = self.cfg.stride if hasattr(self.cfg, 'stride') else self.cfg.window_size // 2\n            start = sample_info['win_idx'] * stride\n            if num_frames <= self.cfg.window_size:\n                start = 0\n\n        if start + self.cfg.window_size > num_frames:\n            start = num_frames - self.cfg.window_size\n\n        if self.cfg.window_size >= num_frames:\n            start = 0\n    \n        end = min(start + self.cfg.window_size, num_frames)\n        features = self.video_keypoints[video_id][start:end]  # Direct slice\n        features = torch.from_numpy(features)\n        labels = None\n\n        actual_len = features.shape[0]\n        if actual_len < self.cfg.window_size:\n            pad_len = self.cfg.window_size - actual_len\n            features = torch.cat([features, torch.full((pad_len, *features.shape[1:]), float('nan'))], dim=0)\n            if labels is not None:\n                labels = torch.cat([labels, torch.zeros(pad_len, *labels.shape[1:])], dim=0)\n                no_action_idx = self.cfg.action_id_map.get('no_action', 38)\n                labels[actual_len:, :, no_action_idx] = 1.0 \n\n        # Create mask\n        mask = torch.ones(self.cfg.window_size)\n        if actual_len < self.cfg.window_size:\n            mask[actual_len:] = 0\n\n        # Apply time reversal if configured\n        if self.cfg.reverse_time:\n            features = torch.flip(features, dims=[0])\n            if labels is not None:\n                labels = torch.flip(labels, dims=[0])\n            mask = torch.flip(mask, dims=[0])\n\n        ego_features_copy = features.clone()\n        features, per_mouse_features, node_feats, edge_feats = self.compute_feature(features)\n        ego_features, ego_per_mouse_features = self.compute_feature_ego(ego_features_copy)\n\n        item = {\n            'input': features,\n            'input_ego': ego_features,\n            'input_mice': per_mouse_features,\n            'input_ego_mice': ego_per_mouse_features,\n            'input_mask': mask,\n            'node_feats': node_feats,\n            'edge_feats': edge_feats,\n            'labels': labels,\n            'video_id': video_id,\n            'start_frame': torch.tensor(start),\n            'num_frames': torch.tensor(num_frames)  # Total video frames\n        }\n\n        # Add behavior_mask (per video)\n        if 'behavior_mask' in video_data:\n            item['behavior_mask'] = torch.from_numpy(video_data['behavior_mask']).float()\n\n        return item\n\n    def compute_egocentric_transform(self, keypoints, bp_to_idx):\n        T, M, B, C = keypoints.shape\n        device = keypoints.device\n        \n        centers = keypoints[:, :, bp_to_idx['body_center']].clone() # [T, M, 2]\n        \n        heads = keypoints[:, :, bp_to_idx['head_center']]\n        tails = keypoints[:, :, bp_to_idx['tail_base']]\n        \n        spines = heads - tails # [T, M, 2]\n\n        spines = torch.nan_to_num(spines, 0.0)\n        \n        spine_norms = torch.norm(spines, dim=-1, keepdim=True)\n        tiny_mask = spine_norms < 1e-6\n        spines = torch.where(tiny_mask, torch.tensor([1.0, 0.0], device=device), spines)\n        spine_norms = torch.where(tiny_mask, torch.tensor(1.0, device=device), spine_norms)\n        \n        cos_theta = spines[..., 0:1] / spine_norms\n        sin_theta = spines[..., 1:2] / spine_norms\n        \n\n        row1 = torch.cat([cos_theta, sin_theta], dim=-1)\n        row2 = torch.cat([-sin_theta, cos_theta], dim=-1)\n        \n        rot_matrices = torch.stack([row1, row2], dim=-2)\n        \n        return centers, rot_matrices\n\n    def compute_feature_ego(self, keypoints):\n        T, M, B, C = keypoints.shape\n        device = keypoints.device\n        \n        # Configuration\n        FPS = 30\n        MAX_VELOCITY = 100.0\n        MAX_ACCEL = 500.0\n        MAX_JERK = 1000.0\n        EPSILON = 1e-6\n        \n        bp_to_idx = self.cfg.MASTER_SKELETON_MAP\n        \n        def safe_divide(a, b, eps=EPSILON):\n            return a / (b + eps)\n        \n        # Helper function for angle computation\n        def compute_angle_3points(p1, p2, p3):\n            v1 = p1 - p2\n            v2 = p3 - p2\n            norm_v1 = torch.norm(v1, dim=-1, keepdim=True)\n            norm_v2 = torch.norm(v2, dim=-1, keepdim=True)\n            cos_angle = torch.sum(v1 * v2, dim=-1) / (norm_v1.squeeze(-1) * norm_v2.squeeze(-1) + EPSILON)\n            return torch.acos(torch.clamp(cos_angle, -1.0 + EPSILON, 1.0 - EPSILON))\n        \n        all_features = []\n        \n        # ============== KINEMATICS (GLOBAL) ==============\n        velocities = torch.diff(keypoints, dim=0, prepend=keypoints[0:1]) * FPS\n        velocities = torch.clamp(velocities, -MAX_VELOCITY, MAX_VELOCITY)\n        accelerations = torch.diff(velocities, dim=0, prepend=velocities[0:1])\n        accelerations = torch.clamp(accelerations, -MAX_ACCEL, MAX_ACCEL)\n        jerks = torch.diff(accelerations, dim=0, prepend=accelerations[0:1])\n        jerks = torch.clamp(jerks, -MAX_JERK, MAX_JERK)\n        \n        speeds_normalized = torch.norm(velocities, dim=-1) / MAX_VELOCITY\n        accel_normalized = torch.norm(accelerations, dim=-1) / MAX_ACCEL\n        jerk_normalized = torch.norm(jerks, dim=-1) / MAX_JERK\n        \n        all_features.append(speeds_normalized.reshape(T, M * B))\n        all_features.append(accel_normalized.reshape(T, M * B))\n        all_features.append(jerk_normalized.reshape(T, M * B))\n        \n        # Angular velocities\n        angular_velocities = torch.zeros(T, M, B).to(device)\n        for m in range(M):\n            for b in range(B):\n                vx, vy = velocities[:, m, b, 0], velocities[:, m, b, 1]\n                angles = torch.atan2(vy, vx)\n                angular_velocities[1:, m, b] = torch.diff(angles)\n                angular_velocities[:, m, b] = (angular_velocities[:, m, b] + torch.pi) % (2 * torch.pi) - torch.pi\n        \n        angular_velocities_normalized = angular_velocities / torch.pi\n        all_features.append(angular_velocities_normalized.reshape(T, M * B))\n        \n        # ============== CENTROID & SHAPE ==============\n        centroid_parts = ['head_center', 'body_center', 'tail_base']\n        centroid_indices = torch.tensor([bp_to_idx[bp] for bp in centroid_parts]).to(device)\n        mouse_centroids = torch.mean(keypoints[:, :, centroid_indices, :], dim=2)\n        \n        centroid_velocity = torch.diff(mouse_centroids, dim=0, prepend=mouse_centroids[0:1]) * FPS\n        centroid_velocity = torch.clamp(centroid_velocity, -MAX_VELOCITY, MAX_VELOCITY)\n        centroid_speed_normalized = torch.norm(centroid_velocity, dim=-1) / MAX_VELOCITY\n        \n        centroid_accel = torch.diff(centroid_velocity, dim=0, prepend=centroid_velocity[0:1])\n        centroid_accel_normalized = torch.norm(torch.clamp(centroid_accel, -MAX_ACCEL, MAX_ACCEL), dim=-1) / MAX_ACCEL\n        \n        all_features.extend([centroid_speed_normalized, centroid_accel_normalized])\n        \n        movement_heading = torch.atan2(centroid_velocity[:, :, 1], centroid_velocity[:, :, 0])\n        movement_heading_normalized = movement_heading / torch.pi\n        all_features.append(movement_heading_normalized)\n        \n        # Shape / Body Configuration\n        nose_pos = keypoints[:, :, bp_to_idx['nose'], :]\n        tail_base_pos = keypoints[:, :, bp_to_idx['tail_base'], :]\n        head_center_pos = keypoints[:, :, bp_to_idx['head_center'], :]\n        body_center_pos = keypoints[:, :, bp_to_idx['body_center'], :]\n        ear_left_pos = keypoints[:, :, bp_to_idx['ear_left'], :]\n        ear_right_pos = keypoints[:, :, bp_to_idx['ear_right'], :]\n        \n        body_length = torch.norm(nose_pos - tail_base_pos, dim=-1)\n        body_length_95 = torch.quantile(body_length[~torch.isnan(body_length)], 0.95) if body_length[~torch.isnan(body_length)].numel() > 0 else 1.0\n        body_length_normalized = torch.clamp(body_length / (body_length_95 + EPSILON), 0, 2)\n        all_features.append(body_length_normalized)\n        \n        body_length_change = torch.diff(body_length, dim=0, prepend=body_length[0:1])\n        body_length_change_normalized = torch.clamp(body_length_change / (body_length_95 * 0.1 + EPSILON), -2, 2)\n        all_features.append(body_length_change_normalized)\n        \n        ear_spread = torch.norm(ear_left_pos - ear_right_pos, dim=-1)\n        ear_spread_95 = torch.quantile(ear_spread[~torch.isnan(ear_spread)], 0.95) if ear_spread[~torch.isnan(ear_spread)].numel() > 0 else 1.0\n        ear_spread_normalized = torch.clamp(ear_spread / (ear_spread_95 + EPSILON), 0, 2)\n        all_features.append(ear_spread_normalized)\n        \n        body_elongation = safe_divide(body_length, ear_spread)\n        body_elongation_normalized = torch.clamp(body_elongation, 0, 10) / 10\n        all_features.append(body_elongation_normalized)\n        \n        # Angles\n        body_curvature = compute_angle_3points(nose_pos, body_center_pos, tail_base_pos)\n        body_curvature_normalized = body_curvature / torch.pi\n        all_features.append(body_curvature_normalized)\n        \n        head_angle = compute_angle_3points(nose_pos, head_center_pos, body_center_pos)\n        head_angle_normalized = head_angle / torch.pi\n        all_features.append(head_angle_normalized)\n        \n        tail_angle = compute_angle_3points(head_center_pos, body_center_pos, tail_base_pos)\n        tail_angle_normalized = tail_angle / torch.pi\n        all_features.append(tail_angle_normalized)\n        \n        body_vector = nose_pos - tail_base_pos\n        body_orientation = torch.atan2(body_vector[:, :, 1], body_vector[:, :, 0])\n        body_orientation_normalized = body_orientation / torch.pi\n        all_features.append(body_orientation_normalized)\n        \n        unwrapped_orientation = torch.from_numpy(np.unwrap(body_orientation.cpu().numpy(), axis=0)).to(device)\n        body_angular_velocity = torch.diff(unwrapped_orientation, dim=0, prepend=unwrapped_orientation[0:1])\n        body_angular_velocity = torch.clamp(body_angular_velocity, -torch.pi, torch.pi) / torch.pi\n        all_features.append(body_angular_velocity)\n        \n        nose_centroid_dist = torch.norm(nose_pos - mouse_centroids, dim=-1)\n        nose_centroid_normalized = torch.clamp(safe_divide(nose_centroid_dist, body_length), 0, 2)\n        tail_centroid_dist = torch.norm(tail_base_pos - mouse_centroids, dim=-1)\n        tail_centroid_normalized = torch.clamp(safe_divide(tail_centroid_dist, body_length), 0, 2)\n        all_features.extend([nose_centroid_normalized, tail_centroid_normalized])\n        \n        nose_velocity = velocities[:, :, bp_to_idx['nose'], :]\n        tail_velocity = velocities[:, :, bp_to_idx['tail_base'], :]\n        nose_rel_speed = torch.norm(nose_velocity - centroid_velocity, dim=-1) / MAX_VELOCITY\n        tail_rel_speed = torch.norm(tail_velocity - centroid_velocity, dim=-1) / MAX_VELOCITY\n        all_features.extend([nose_rel_speed, tail_rel_speed])\n        \n        # ============== EGOCENTRIC ALIGNMENT (THE FIX) ==============\n        ego_centers, ego_rot_mats = self.compute_egocentric_transform(keypoints, bp_to_idx)\n        \n        typical_mouse_size = torch.median(body_length[~torch.isnan(body_length)]) if body_length[~torch.isnan(body_length)].numel() > 0 else 1.0\n\n        # ============== INTER-MOUSE FEATURES ==============\n        inter_features_list = []\n        \n        # Per-mouse lists\n        mouse_features = [[] for _ in range(M)]\n        \n        # 1. Add Intra-mouse features\n        bp_specific = [speeds_normalized, accel_normalized, jerk_normalized, angular_velocities_normalized]\n        for feat in bp_specific:\n            for m in range(M):\n                mouse_features[m].append(feat[:, m, :].reshape(T, B))\n        \n        mouse_specific = [\n            centroid_speed_normalized, centroid_accel_normalized, movement_heading_normalized,\n            body_length_normalized, body_length_change_normalized, ear_spread_normalized, body_elongation_normalized,\n            body_curvature_normalized, head_angle_normalized, tail_angle_normalized, body_orientation_normalized, body_angular_velocity,\n            nose_centroid_normalized, tail_centroid_normalized, nose_rel_speed, tail_rel_speed\n        ]\n        for feat in mouse_specific:\n            for m in range(M):\n                mouse_features[m].append(feat[:, m].unsqueeze(-1))\n        \n        # 2. Add Egocentric Interaction Features\n        for i in range(M): # Agent\n            R_i = ego_rot_mats[:, i] # [T, 2, 2]\n            Center_i = ego_centers[:, i] # [T, 2]\n            \n            for j in range(M): # Target\n                if i == j:\n                    continue\n                \n                # --- Basic Distances (Existing) ---\n                centroid_i = mouse_centroids[:, i, :]\n                centroid_j = mouse_centroids[:, j, :]\n                dist_centroids = torch.norm(centroid_j - centroid_i, dim=-1)\n                dist_centroids_normalized = torch.clamp(dist_centroids / (typical_mouse_size * 5 + EPSILON), 0, 2)\n                mouse_features[i].append(dist_centroids_normalized.unsqueeze(-1))\n                inter_features_list.append(dist_centroids_normalized.unsqueeze(-1))\n                \n                dist_change = torch.diff(dist_centroids, dim=0, prepend=dist_centroids[0:1])\n                dist_change_normalized = torch.clamp(dist_change / (typical_mouse_size + EPSILON), -2, 2) / 2\n                mouse_features[i].append(dist_change_normalized.unsqueeze(-1))\n                inter_features_list.append(dist_change_normalized.unsqueeze(-1))\n                \n                nose_i = nose_pos[:, i, :]\n                nose_j = nose_pos[:, j, :]\n                tail_j = tail_base_pos[:, j, :]\n                \n                nose_nose_dist = torch.norm(nose_i - nose_j, dim=-1)\n                nose_nose_normalized = torch.clamp(nose_nose_dist / (typical_mouse_size * 3 + EPSILON), 0, 2)\n                mouse_features[i].append(nose_nose_normalized.unsqueeze(-1))\n                inter_features_list.append(nose_nose_normalized.unsqueeze(-1))\n                \n                nose_tail_dist = torch.norm(nose_i - tail_j, dim=-1)\n                nose_tail_normalized = torch.clamp(nose_tail_dist / (typical_mouse_size * 3 + EPSILON), 0, 2)\n                mouse_features[i].append(nose_tail_normalized.unsqueeze(-1))\n                inter_features_list.append(nose_tail_normalized.unsqueeze(-1))\n                \n                rel_speed = (centroid_speed_normalized[:, i] - centroid_speed_normalized[:, j])\n                mouse_features[i].append(rel_speed.unsqueeze(-1))\n                inter_features_list.append(rel_speed.unsqueeze(-1))\n                \n                vec_to_j = centroid_j - centroid_i\n                heading_vec_i = body_vector[:, i, :]\n                dot_prod = torch.sum(heading_vec_i * vec_to_j, dim=-1)\n                norm_prod = torch.norm(heading_vec_i, dim=-1) * torch.norm(vec_to_j, dim=-1)\n                approach_angle = torch.acos(torch.clamp(safe_divide(dot_prod, norm_prod), -1.0 + EPSILON, 1.0 - EPSILON))\n                approach_angle_normalized = approach_angle / torch.pi\n                mouse_features[i].append(approach_angle_normalized.unsqueeze(-1))\n                inter_features_list.append(approach_angle_normalized.unsqueeze(-1))\n                \n                # --- NEW EGOCENTRIC FEATURES (High SOTA Impact) ---\n                target_pts = {\n                    'nose': nose_pos[:, j, :],\n                    'body': mouse_centroids[:, j, :],\n                    'tail': tail_base_pos[:, j, :]\n                }\n                \n                for pt_name, pt_tensor in target_pts.items():\n                    rel_vec = pt_tensor - Center_i # [T, 2]\n                    \n                    rel_vec_ego = torch.matmul(R_i, rel_vec.unsqueeze(-1)).squeeze(-1) # [T, 2]\n                    \n                    norm_scale = typical_mouse_size * 5.0 + EPSILON\n                    \n                    ego_x = torch.clamp(rel_vec_ego[:, 0] / norm_scale, -2.0, 2.0)\n                    ego_y = torch.clamp(rel_vec_ego[:, 1] / norm_scale, -2.0, 2.0)\n                    \n                    # Add to features\n                    mouse_features[i].append(ego_x.unsqueeze(-1))\n                    mouse_features[i].append(ego_y.unsqueeze(-1))\n                    \n                    # Also add to flattened inter_features_list for global stream\n                    inter_features_list.append(ego_x.unsqueeze(-1))\n                    inter_features_list.append(ego_y.unsqueeze(-1))\n\n        if inter_features_list:\n            inter_features = torch.cat(inter_features_list, dim=-1)\n            all_features.append(inter_features)\n        \n        per_mouse_feats = [torch.cat(feats, dim=-1) for feats in mouse_features]\n        \n        features = torch.cat(all_features, dim=-1)\n        per_mouse_feats = torch.stack(per_mouse_feats, dim=1)\n\n        features = torch.clamp(features, -10, 10)\n        per_mouse_feats = torch.clamp(per_mouse_feats, -10, 10)\n        \n        features = torch.nan_to_num(features, nan=0.0, posinf=10.0, neginf=-10.0)\n        per_mouse_feats = torch.nan_to_num(per_mouse_feats, nan=0.0, posinf=10.0, neginf=-10.0)\n        \n        return features, per_mouse_feats\n\n    def compute_feature(self, keypoints):\n            T, M, B, C = keypoints.shape\n            device = keypoints.device\n            \n            # Configuration\n            FPS = 30  \n            WINDOW_SIZES = [3, 5, 10, 20]  # For rolling statistics\n            \n            # Clipping thresholds to prevent explosions\n            MAX_VELOCITY = 100.0  # cm/s\n            MAX_ACCEL = 500.0     # cm/s^2\n            MAX_JERK = 1000.0     # cm/s^3\n            EPSILON = 1e-6\n            \n            # Helper function for safe division\n            def safe_divide(a, b, eps=EPSILON):\n                return a / (b + eps)\n            \n            # Helper function for angle computation\n            def compute_angle_3points(p1, p2, p3):\n                \"\"\"Compute angle at p2 formed by p1-p2-p3\"\"\"\n                v1 = p1 - p2\n                v2 = p3 - p2\n                norm_v1 = torch.norm(v1, dim=-1, keepdim=True)\n                norm_v2 = torch.norm(v2, dim=-1, keepdim=True)\n                cos_angle = torch.sum(v1 * v2, dim=-1) / (norm_v1.squeeze(-1) * norm_v2.squeeze(-1) + EPSILON)\n                return torch.acos(torch.clamp(cos_angle, -1.0 + EPSILON, 1.0 - EPSILON))\n            \n            # Helper for rolling statistics\n            def rolling_window_torch(tensor, window_size, operation='mean'):\n                \"\"\"Apply rolling window operation\"\"\"\n                result = torch.zeros_like(tensor)\n                pad = window_size // 2\n                for i in range(tensor.shape[0]):\n                    start = max(0, i - pad)\n                    end = min(tensor.shape[0], i + pad + 1)\n                    if operation == 'mean':\n                        result[i] = torch.mean(tensor[start:end], dim=0)\n                    elif operation == 'std':\n                        if end - start > 1:\n                            result[i] = torch.std(tensor[start:end], dim=0)\n                        else:\n                            result[i] = 0\n                return result\n            \n            all_features = []\n            bp_to_idx = self.cfg.MASTER_SKELETON_MAP\n            \n            # ============== KINEMATICS (NORMALIZED) ==============\n            \n            # Velocities with clipping\n            velocities = torch.diff(keypoints, dim=0, prepend=keypoints[0:1]) * FPS\n            velocities = torch.clamp(velocities, -MAX_VELOCITY, MAX_VELOCITY)\n            \n            # Accelerations with clipping\n            accelerations = torch.diff(velocities, dim=0, prepend=velocities[0:1])\n            accelerations = torch.clamp(accelerations, -MAX_ACCEL, MAX_ACCEL)\n            \n            # Jerks with clipping\n            jerks = torch.diff(accelerations, dim=0, prepend=accelerations[0:1])\n            jerks = torch.clamp(jerks, -MAX_JERK, MAX_JERK)\n            \n            # Speed and acceleration magnitudes (normalized by max values)\n            speeds = torch.norm(velocities, dim=-1)  # [T, M, B]\n            speeds_normalized = speeds / MAX_VELOCITY  # Normalize to [0, 1]\n            \n            accel_mags = torch.norm(accelerations, dim=-1)  # [T, M, B]\n            accel_normalized = accel_mags / MAX_ACCEL\n            \n            jerk_mags = torch.norm(jerks, dim=-1)  # [T, M, B]\n            jerk_normalized = jerk_mags / MAX_JERK\n            \n            all_features.append(speeds_normalized.reshape(T, M * B))\n            all_features.append(accel_normalized.reshape(T, M * B))\n            all_features.append(jerk_normalized.reshape(T, M * B))\n            \n            # Angular velocities (already bounded [-pi, pi] naturally)\n            angular_velocities = torch.zeros(T, M, B).to(device)\n            for m in range(M):\n                for b in range(B):\n                    vx, vy = velocities[:, m, b, 0], velocities[:, m, b, 1]\n                    angles = torch.atan2(vy, vx)\n                    angular_velocities[1:, m, b] = torch.diff(angles)\n                    # Wrap to [-pi, pi]\n                    angular_velocities[:, m, b] = (angular_velocities[:, m, b] + torch.pi) % (2 * torch.pi) - torch.pi\n            \n            # Normalize angular velocities to [-1, 1]\n            angular_velocities_normalized = angular_velocities / torch.pi\n            all_features.append(angular_velocities_normalized.reshape(T, M * B))\n            \n            # ============== CENTER OF MASS FEATURES ==============\n            \n            # Centroid using key body parts\n            centroid_parts = ['head_center', 'body_center', 'tail_base']\n            centroid_indices = torch.tensor([bp_to_idx[bp] for bp in centroid_parts]).to(device)\n            mouse_centroids = torch.mean(keypoints[:, :, centroid_indices, :], dim=2)  # [T, M, 2]\n            \n            # Centroid kinematics (normalized)\n            centroid_velocity = torch.diff(mouse_centroids, dim=0, prepend=mouse_centroids[0:1]) * FPS\n            centroid_velocity = torch.clamp(centroid_velocity, -MAX_VELOCITY, MAX_VELOCITY)\n            centroid_speed = torch.norm(centroid_velocity, dim=-1)  # [T, M]\n            centroid_speed_normalized = centroid_speed / MAX_VELOCITY\n            \n            centroid_accel = torch.diff(centroid_velocity, dim=0, prepend=centroid_velocity[0:1])\n            centroid_accel = torch.clamp(centroid_accel, -MAX_ACCEL, MAX_ACCEL)\n            centroid_accel_mag = torch.norm(centroid_accel, dim=-1)\n            centroid_accel_normalized = centroid_accel_mag / MAX_ACCEL\n            \n            all_features.extend([centroid_speed_normalized, centroid_accel_normalized])\n            \n            # Movement direction (already bounded [-pi, pi])\n            movement_heading = torch.atan2(centroid_velocity[:, :, 1], centroid_velocity[:, :, 0])\n            movement_heading_normalized = movement_heading / torch.pi  # [-1, 1]\n            all_features.append(movement_heading_normalized)\n            \n            # ============== BODY CONFIGURATION (NORMALIZED) ==============\n            \n            # Key positions\n            nose_pos = keypoints[:, :, bp_to_idx['nose'], :]\n            tail_base_pos = keypoints[:, :, bp_to_idx['tail_base'], :]\n            head_center_pos = keypoints[:, :, bp_to_idx['head_center'], :]\n            body_center_pos = keypoints[:, :, bp_to_idx['body_center'], :]\n            ear_left_pos = keypoints[:, :, bp_to_idx['ear_left'], :]\n            ear_right_pos = keypoints[:, :, bp_to_idx['ear_right'], :]\n            \n            # Body length (normalize by percentile or known max)\n            body_length = torch.norm(nose_pos - tail_base_pos, dim=-1)  # [T, M]\n            # Use 95th percentile for robust normalization\n            body_length_95 = torch.quantile(body_length[~torch.isnan(body_length)], 0.95) if body_length[~torch.isnan(body_length)].numel() > 0 else 1.0\n            body_length_normalized = torch.clamp(body_length / (body_length_95 + EPSILON), 0, 2)  # Allow some values > 1\n            all_features.append(body_length_normalized)\n            \n            # Body length rate of change (normalized)\n            body_length_change = torch.diff(body_length, dim=0, prepend=body_length[0:1])\n            body_length_change_normalized = torch.clamp(body_length_change / (body_length_95 * 0.1 + EPSILON), -2, 2)\n            all_features.append(body_length_change_normalized)\n            \n            # Ear spread (normalized)\n            ear_spread = torch.norm(ear_left_pos - ear_right_pos, dim=-1)  # [T, M]\n            ear_spread_95 = torch.quantile(ear_spread[~torch.isnan(ear_spread)], 0.95) if ear_spread[~torch.isnan(ear_spread)].numel() > 0 else 1.0\n            ear_spread_normalized = torch.clamp(ear_spread / (ear_spread_95 + EPSILON), 0, 2)\n            all_features.append(ear_spread_normalized)\n            \n            # Body elongation ratio (already a ratio, just clamp)\n            body_elongation = safe_divide(body_length, ear_spread)\n            body_elongation_normalized = torch.clamp(body_elongation, 0, 10) / 10  # Typical range 1-5\n            all_features.append(body_elongation_normalized)\n            \n            # ============== BODY ANGLES (ALREADY NORMALIZED) ==============\n            \n            # All angles are in radians, normalize to [-1, 1]\n            body_curvature = compute_angle_3points(nose_pos, body_center_pos, tail_base_pos)  # [T, M]\n            body_curvature_normalized = body_curvature / torch.pi\n            all_features.append(body_curvature_normalized)\n            \n            head_angle = compute_angle_3points(nose_pos, head_center_pos, body_center_pos)  # [T, M]\n            head_angle_normalized = head_angle / torch.pi\n            all_features.append(head_angle_normalized)\n            \n            tail_angle = compute_angle_3points(head_center_pos, body_center_pos, tail_base_pos)  # [T, M]\n            tail_angle_normalized = tail_angle / torch.pi\n            all_features.append(tail_angle_normalized)\n            \n            # Body orientation angle\n            body_vector = nose_pos - tail_base_pos\n            body_orientation = torch.atan2(body_vector[:, :, 1], body_vector[:, :, 0])  # [T, M]\n            body_orientation_normalized = body_orientation / torch.pi\n            all_features.append(body_orientation_normalized)\n            \n            # Angular velocity of body orientation\n            unwrapped_orientation = torch.from_numpy(np.unwrap(body_orientation.cpu().numpy(), axis=0)).to(device)\n            body_angular_velocity = torch.diff(unwrapped_orientation, dim=0, prepend=unwrapped_orientation[0:1])\n            body_angular_velocity = torch.clamp(body_angular_velocity, -torch.pi, torch.pi) / torch.pi\n            all_features.append(body_angular_velocity)\n            \n            # ============== EXTREMITY FEATURES (NORMALIZED) ==============\n            \n            # Normalize distances by body length\n            nose_centroid_dist = torch.norm(nose_pos - mouse_centroids, dim=-1)  # [T, M]\n            nose_centroid_normalized = safe_divide(nose_centroid_dist, body_length)\n            nose_centroid_normalized = torch.clamp(nose_centroid_normalized, 0, 2)\n            \n            tail_centroid_dist = torch.norm(tail_base_pos - mouse_centroids, dim=-1)  # [T, M]\n            tail_centroid_normalized = safe_divide(tail_centroid_dist, body_length)\n            tail_centroid_normalized = torch.clamp(tail_centroid_normalized, 0, 2)\n            \n            all_features.extend([nose_centroid_normalized, tail_centroid_normalized])\n            \n            # Relative velocities (normalized)\n            nose_velocity = velocities[:, :, bp_to_idx['nose'], :]\n            tail_velocity = velocities[:, :, bp_to_idx['tail_base'], :]\n            nose_rel_velocity = nose_velocity - centroid_velocity\n            tail_rel_velocity = tail_velocity - centroid_velocity\n            nose_rel_speed = torch.norm(nose_rel_velocity, dim=-1) / MAX_VELOCITY\n            tail_rel_speed = torch.norm(tail_rel_velocity, dim=-1) / MAX_VELOCITY\n            all_features.extend([nose_rel_speed, tail_rel_speed])\n                    \n            # ============== INTER-MOUSE FEATURES (NORMALIZED) ==============\n            \n            inter_features_list = []\n            \n            # Get typical mouse size for normalization\n            typical_mouse_size = torch.median(body_length[~torch.isnan(body_length)]) if body_length[~torch.isnan(body_length)].numel() > 0 else 1.0\n            \n\n            ######### THIS IS PER MOUSE ADAPTATION\n            # Initialize per-mouse feature lists\n            mouse_features = [[] for _ in range(M)]\n            \n            # Add bodypart-specific features per mouse\n            bp_specific = [speeds_normalized, accel_normalized, jerk_normalized, angular_velocities_normalized]\n            for feat in bp_specific:  # each [T, M, B]\n                for m in range(M):\n                    mouse_features[m].append(feat[:, m, :].reshape(T, B))  # [T, 6]\n            \n            # Add mouse-specific features per mouse\n            mouse_specific = [\n                centroid_speed_normalized, centroid_accel_normalized, movement_heading_normalized,\n                body_length_normalized, body_length_change_normalized, ear_spread_normalized, body_elongation_normalized,\n                body_curvature_normalized, head_angle_normalized, tail_angle_normalized, body_orientation_normalized, body_angular_velocity,\n                nose_centroid_normalized, tail_centroid_normalized, nose_rel_speed, tail_rel_speed\n            ]\n            for feat in mouse_specific:  # each [T, M]\n                for m in range(M):\n                    mouse_features[m].append(feat[:, m].unsqueeze(-1))  # [T, 1]\n\n            pure_node_feats_list = [torch.cat(feats, dim=-1) for feats in mouse_features]\n            pure_node_feats = torch.stack(pure_node_feats_list, dim=1) # [T, M, Node_Dim]\n            pure_node_feats = torch.clamp(pure_node_feats, -10, 10)\n            pure_node_feats = torch.nan_to_num(pure_node_feats, nan=0.0, posinf=10.0, neginf=-10.0)\n\n            edge_attrs_matrix = torch.zeros(T, M, M, 8).to(device)\n            \n            for i in range(M):\n                for j in range(M):\n                    if i == j:\n                        continue\n                    \n                    # Distances normalized by typical mouse size\n                    centroid_i = mouse_centroids[:, i, :]\n                    centroid_j = mouse_centroids[:, j, :]\n                    dist_centroids = torch.norm(centroid_j - centroid_i, dim=-1)\n                    dist_centroids_normalized = torch.clamp(dist_centroids / (typical_mouse_size * 5 + EPSILON), 0, 2)\n                    \n                    mouse_features[i].append(dist_centroids_normalized.unsqueeze(-1))  # From i's perspective\n                    inter_features_list.append(dist_centroids_normalized.unsqueeze(-1))\n                    \n                    # Rate of distance change (normalized)\n                    dist_change = torch.diff(dist_centroids, dim=0, prepend=dist_centroids[0:1])\n                    dist_change_normalized = torch.clamp(dist_change / (typical_mouse_size + EPSILON), -2, 2) / 2\n                    \n                    mouse_features[i].append(dist_change_normalized.unsqueeze(-1))\n                    inter_features_list.append(dist_change_normalized.unsqueeze(-1))\n                    \n                    nose_i = nose_pos[:, i, :]\n                    nose_j = nose_pos[:, j, :]\n                    tail_j = tail_base_pos[:, j, :]\n                    \n                    nose_nose_dist = torch.norm(nose_i - nose_j, dim=-1)\n                    nose_nose_normalized = torch.clamp(nose_nose_dist / (typical_mouse_size * 3 + EPSILON), 0, 2)\n                    mouse_features[i].append(nose_nose_normalized.unsqueeze(-1))\n                    \n                    nose_tail_dist = torch.norm(nose_i - tail_j, dim=-1)\n                    nose_tail_normalized = torch.clamp(nose_tail_dist / (typical_mouse_size * 3 + EPSILON), 0, 2)\n                    mouse_features[i].append(nose_tail_normalized.unsqueeze(-1))\n                    \n                    inter_features_list.extend([\n                        nose_nose_normalized.unsqueeze(-1),\n                        nose_tail_normalized.unsqueeze(-1)\n                    ])\n                    \n                    rel_speed = (centroid_speed_normalized[:, i] - centroid_speed_normalized[:, j])\n                    mouse_features[i].append(rel_speed.unsqueeze(-1))\n                    inter_features_list.append(rel_speed.unsqueeze(-1))\n                    \n                    # Approach angle (normalized)\n                    vec_to_j = centroid_j - centroid_i\n                    heading_vec_i = body_vector[:, i, :]\n                    dot_prod = torch.sum(heading_vec_i * vec_to_j, dim=-1)\n                    norm_prod = torch.norm(heading_vec_i, dim=-1) * torch.norm(vec_to_j, dim=-1)\n                    approach_angle = torch.acos(torch.clamp(safe_divide(dot_prod, norm_prod), -1.0 + EPSILON, 1.0 - EPSILON))\n                    approach_angle_normalized = approach_angle / torch.pi\n                    \n                    inter_features_list.append(approach_angle_normalized.unsqueeze(-1))\n                    mouse_features[i].append(approach_angle_normalized.unsqueeze(-1))\n\n                    heading_vec_j = body_vector[:, j, :]\n                    dot_align = torch.sum(heading_vec_i * heading_vec_j, dim=-1)\n                    norm_h_i = torch.norm(heading_vec_i, dim=-1) + EPSILON\n                    norm_h_j = torch.norm(heading_vec_j, dim=-1) + EPSILON\n                    alignment = torch.clamp(dot_align / (norm_h_i * norm_h_j), -1.0, 1.0)\n                    \n                    vec_to_i = -vec_to_j # Vector from J to I\n                    dot_aspect = torch.sum(heading_vec_j * vec_to_i, dim=-1)\n                    norm_vec_i = torch.norm(vec_to_i, dim=-1) + EPSILON\n                    target_aspect = torch.acos(torch.clamp(dot_aspect / (norm_h_j * norm_vec_i), -1.0 + EPSILON, 1.0 - EPSILON))\n                    target_aspect_norm = target_aspect / torch.pi\n\n                    edge_attrs_matrix[:, i, j, 0] = dist_centroids_normalized\n                    edge_attrs_matrix[:, i, j, 1] = dist_change_normalized\n                    edge_attrs_matrix[:, i, j, 2] = nose_nose_normalized\n                    edge_attrs_matrix[:, i, j, 3] = nose_tail_normalized\n                    edge_attrs_matrix[:, i, j, 4] = rel_speed\n                    edge_attrs_matrix[:, i, j, 5] = approach_angle_normalized\n                    edge_attrs_matrix[:, i, j, 6] = alignment\n                    edge_attrs_matrix[:, i, j, 7] = target_aspect_norm\n                    \n            if inter_features_list:\n                inter_features = torch.cat(inter_features_list, dim=-1)\n                all_features.append(inter_features)\n            \n            per_mouse_feats = [torch.cat(feats, dim=-1) for feats in mouse_features]  # List of [T, 58]\n            # ============== CONCATENATE AND FINAL NORMALIZATION ==============\n            \n            features = torch.cat(all_features, dim=-1)\n            per_mouse_feats = torch.stack(per_mouse_feats, dim=1)  # [T, M, 58]\n\n            features = torch.clamp(features, -10, 10)\n            per_mouse_feats = torch.clamp(per_mouse_feats, -10, 10)\n            \n            # Handle NaN/Inf\n            features = torch.nan_to_num(features, nan=0.0, posinf=10.0, neginf=-10.0)\n            per_mouse_feats = torch.nan_to_num(per_mouse_feats, nan=0.0, posinf=10.0, neginf=-10.0)\n            edge_attrs_matrix = torch.nan_to_num(edge_attrs_matrix, nan=0.0, posinf=10.0, neginf=-10.0)\n\n            return features, per_mouse_feats, pure_node_feats, edge_attrs_matrix\n                \n\n\ndef collate_fn(batch):\n    \"\"\"Custom collate function for batching\"\"\"\n    inputs = torch.stack([item['input'] for item in batch])\n    inputs_ego = torch.stack([item['input_ego'] for item in batch]) if 'input_ego' in batch[0] else None\n    inputs_mice = torch.stack([item['input_mice'] for item in batch]) if 'input_mice' in batch[0] else None\n    inputs_ego_mice = torch.stack([item['input_ego_mice'] for item in batch]) if 'input_ego_mice' in batch[0] else None\n    node_feats = torch.stack([item['node_feats'] for item in batch]) if 'node_feats' in batch[0] else None\n    edge_feats = torch.stack([item['edge_feats'] for item in batch]) if 'edge_feats' in batch[0] else None\n    masks = torch.stack([item['input_mask'] for item in batch])\n    behavior_masks = torch.stack([item['behavior_mask'] for item in batch]) if 'behavior_mask' in batch[0] else None\n\n    labels = torch.stack([item['labels'] for item in batch]) if batch[0]['labels'] is not None else None\n\n    video_ids = [item['video_id'] for item in batch]\n    start_frames = torch.stack([item['start_frame'] for item in batch])\n    num_frames = torch.stack([item['num_frames'] for item in batch])\n\n    collated = {\n        'input': inputs,\n        'input_ego': inputs_ego,\n        'input_mice': inputs_mice,\n        'input_ego_mice': inputs_ego_mice,\n        'node_feats': node_feats,\n        'edge_feats': edge_feats,\n        'input_mask': masks,\n        'labels': labels,\n        'video_id': video_ids,\n        'start_frame': start_frames,\n        'num_frames': num_frames\n    }\n    if behavior_masks is not None:\n        collated['behavior_mask'] = behavior_masks\n\n    return collated\n    \n# For compatibility\ntr_collate_fn = collate_fn\nval_collate_fn = collate_fn","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T13:06:49.123290Z","iopub.execute_input":"2025-12-18T13:06:49.123558Z","iopub.status.idle":"2025-12-18T13:06:49.210553Z","shell.execute_reply.started":"2025-12-18T13:06:49.123519Z","shell.execute_reply":"2025-12-18T13:06:49.209918Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### POST PROCESS","metadata":{}},{"cell_type":"code","source":"def process_one_video_preds_filters(cfg, video_preds, video_id, lab_id, num_frames,\n                                       mouse_map, action_map, num_pairs, num_actions,\n                                       min_duration, smooth_sigma=0.0, conf_thresh=0.0,\n                                       max_gap=0, use_median_filter=False):\n\n\n    if smooth_sigma > 0:\n        smoothed_preds = np.zeros_like(video_preds)\n        for pair in range(num_pairs):\n            for act in range(num_actions):\n                smoothed_preds[:, pair, act] = gaussian_filter1d(video_preds[:, pair, act], sigma=smooth_sigma)\n        video_preds = smoothed_preds / np.sum(smoothed_preds, axis=2, keepdims=True)  # Re-normalize\n\n    # Get best actions\n    best_action_indices = np.argmax(video_preds, axis=2)  # [num_frames, num_pairs]\n    best_probs = np.max(video_preds, axis=2)  # [num_frames, num_pairs] for confidence\n\n    # Optional: Median filter on labels (after argmax)\n    if use_median_filter:\n        for pair in range(num_pairs):\n            best_action_indices[:, pair] = median_filter(best_action_indices[:, pair], size=5)  # Window size tunable\n\n    video_predictions = []\n    num_mice = int(np.sqrt(num_pairs))\n\n    for pair_idx in range(num_pairs):\n        agent_idx = pair_idx // num_mice\n        target_idx = pair_idx % num_mice\n        agent_id = mouse_map[agent_idx]\n        target_id = mouse_map[target_idx] if agent_idx != target_idx else 'self'\n\n        for action_idx in range(num_actions):\n            action = action_map[action_idx]\n            if action == 'no_action':\n                continue\n            if (agent_idx == target_idx and action in cfg.social_only_actions) or \\\n               (agent_idx != target_idx and action in cfg.self_only_actions):\n                continue\n\n            behavior_frames = (best_action_indices[:, pair_idx] == action_idx)\n\n            # Find contiguous segments\n            padded = np.concatenate(([False], behavior_frames, [False]))\n            diff = np.diff(padded.astype(int))\n            starts = np.where(diff > 0)[0]\n            ends = np.where(diff < 0)[0] - 1\n\n            # Step 2: Merge gaps (for same action)\n            if len(starts) > 1 and max_gap > 0:\n                merged_starts, merged_ends = [starts[0]], [ends[0]]\n                for i in range(1, len(starts)):\n                    if starts[i] - merged_ends[-1] - 1 <= max_gap:\n                        merged_ends[-1] = ends[i]  # Merge\n                    else:\n                        merged_starts.append(starts[i])\n                        merged_ends.append(ends[i])\n                starts, ends = np.array(merged_starts), np.array(merged_ends)\n\n            # Filter by min_duration and confidence\n            durations = ends - starts + 1\n            valid = durations >= min_duration\n            for start, end in zip(starts[valid], ends[valid]):\n                seg_probs = video_preds[start:end+1, pair_idx, action_idx]  # Use action-specific probs\n                mean_conf = np.mean(seg_probs)\n                if mean_conf < conf_thresh:\n                    continue\n                video_predictions.append({\n                    'video_id': video_id,\n                    'agent_id': agent_id,\n                    'target_id': target_id,\n                    'action': action,\n                    'start_frame': start,\n                    'stop_frame': end,\n                    'confidence': mean_conf\n                })\n\n    if video_predictions:\n        df = pd.DataFrame(video_predictions)\n        resolved_predictions = []\n        for (agent_id, target_id), group in df.groupby(['agent_id', 'target_id']):\n            segments = group.to_dict('records')\n            segments.sort(key=lambda x: x['start_frame'])\n            # Check for overlap\n            overlap_found = False\n            for k in range(1, len(segments)):\n                if segments[k]['start_frame'] <= segments[k-1]['stop_frame']:\n                    overlap_found = True\n                    break\n            if overlap_found:\n                print(f\"Warning: Overlaps found in {video_id} for pair {agent_id}-{target_id}. Resolving with NMS.\")\n                # NMS: Sort by descending confidence, greedily select non-overlapping\n                segments.sort(key=lambda x: -x['confidence'])\n                selected = []\n                for seg in segments:\n                    if all(max(seg['start_frame'], sel['start_frame']) > min(seg['stop_frame'], sel['stop_frame']) for sel in selected):\n                        selected.append(seg)\n                resolved_predictions.extend(selected)\n            else:\n                resolved_predictions.extend(segments)\n        if resolved_predictions:\n            return pd.DataFrame(resolved_predictions)\n    return None\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T13:06:49.211306Z","iopub.execute_input":"2025-12-18T13:06:49.211586Z","iopub.status.idle":"2025-12-18T13:06:49.236037Z","shell.execute_reply.started":"2025-12-18T13:06:49.211556Z","shell.execute_reply":"2025-12-18T13:06:49.235413Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### TRAIN","metadata":{}},{"cell_type":"code","source":"import os\nimport glob\nimport gc\nfrom copy import copy\nimport numpy as np\nimport pandas as pd\nimport importlib\nimport sys\nfrom tqdm import tqdm\nimport argparse\nimport torch\nfrom torch.amp import autocast, GradScaler\nimport json\nfrom joblib import Parallel, delayed\nimport ast\nfrom time import time\nfrom torch.utils.data import DataLoader\nimport warnings\nimport seaborn as sns\nimport matplotlib.pyplot as plt\n\nimport model_1\nimport model_gnn\nimport model_inter\nimport model_238 as model_238_code\nimport model_242 as model_242_code\nimport model_243 as model_243_code\nimport model_244 as model_244_code\n\nwarnings.filterwarnings('ignore')\n\nstart_time = time()\n\ncfg.device = 'cuda' if torch.cuda.is_available() else 'cpu'\nprint(f\"Using device: {cfg.device}\")\n\nwindow_size = cfg.window_size\ncfg.stride = window_size // 2  # 50% overlap\nmin_duration = 5 # Minimum frames for a behavior\n\ncfg_233 = copy(cfg)\ncfg_233.reverse_time = False\ncfg_233.use_gnn = False\ncfg_233.use_bn = True\ncfg_233.cnn_extractor = True\ncfg_233.encoder_config.input_dim=256\ncfg_233.encoder_config.encoder_dim=256\ncfg_233.encoder_config.num_layers=4\ncfg_233.encoder_config.num_attention_heads=4\ncfg_233.per_mouse_feature_dim = cfg_233.feature_dim // 4\n\n\nmodel_233 = model_gnn.NetGNN(cfg_233).to(cfg_233.device)\nmodel_233 = torch.compile(model_233)\nweights_path_233 = \"/kaggle/input/mabe-weights/MAB-233_best.pth\"\nweights_233 = torch.load(weights_path_233, map_location=cfg_233.device, weights_only=False)#['model_state_dict']\nmodel_233.load_state_dict(weights_233)\nmodel_233.eval()\n\n\ncfg_238 = copy(cfg)\ncfg_238.cnn_extractor = True\ncfg_238.use_gnn = True\ncfg_238.use_bn = True\ncfg_238.encoder_config.input_dim=256\ncfg_238.encoder_config.encoder_dim=256\ncfg_238.encoder_config.num_layers=4\ncfg_238.encoder_config.num_attention_heads=4\ncfg_238.per_mouse_feature_dim = cfg_238.feature_dim // 4\n\nmodel_238 = model_238_code.Net(cfg_238).to(cfg_238.device)\nmodel_238 = torch.compile(model_238)\nweights_path_238 = \"/kaggle/input/mabe-weights/MAB-238_best.pth\"\nweights_238 = torch.load(weights_path_238, map_location=cfg_238.device, weights_only=False)#['model_state_dict']\nmodel_238.load_state_dict(weights_238)\nmodel_238.eval()\n\nmodel_289 = model_238_code.Net(cfg_238).to(cfg_238.device)\nmodel_289 = torch.compile(model_289)\nweights_path_289 = \"/kaggle/input/mabe-weights/MAB-289_best.pth\"\nweights_289 = torch.load(weights_path_289, map_location=cfg_238.device, weights_only=False)#['model_state_dict']\nmodel_289.load_state_dict(weights_289)\nmodel_289.eval()\n\nmodel_290 = model_238_code.Net(cfg_238).to(cfg_238.device)\nmodel_290 = torch.compile(model_290)\nweights_path_290 = \"/kaggle/input/mabe-weights/MAB-290_best.pth\"\nweights_290 = torch.load(weights_path_290, map_location=cfg_238.device, weights_only=False)#['model_state_dict']\nmodel_290.load_state_dict(weights_290)\nmodel_290.eval()\n\nmodel_293 = model_238_code.Net(cfg_238).to(cfg_238.device)\nmodel_293 = torch.compile(model_293)\nweights_path_293 = \"/kaggle/input/mabe-weights/MAB-293_best.pth\"\nweights_293 = torch.load(weights_path_293, map_location=cfg_238.device, weights_only=False)#['model_state_dict']\nmodel_293.load_state_dict(weights_293)\nmodel_293.eval()\n\nmodel_292 = model_238_code.Net(cfg_238).to(cfg_238.device)\nmodel_292 = torch.compile(model_292)\nweights_path_292 = \"/kaggle/input/mabe-weights/MAB-292_best.pth\"\nweights_292 = torch.load(weights_path_292, map_location=cfg_238.device, weights_only=False)#['model_state_dict']\nmodel_292.load_state_dict(weights_292)\nmodel_292.eval()\n\nmodel_245 = model_238_code.Net(cfg_238).to(cfg_238.device)\nmodel_245 = torch.compile(model_245)\nweights_path_245 = \"/kaggle/input/mabe-weights/MAB-245_best.pth\"\nweights_245 = torch.load(weights_path_245, map_location=cfg_238.device, weights_only=False)#['model_state_dict']\nmodel_245.load_state_dict(weights_245)\nmodel_245.eval()\n\nmodel_264 = model_238_code.Net(cfg_238).to(cfg_238.device)\nmodel_264 = torch.compile(model_264)\nweights_path_264 = \"/kaggle/input/mabe-weights/MAB-264_best.pth\"\nweights_264 = torch.load(weights_path_264, map_location=cfg_238.device, weights_only=False)#['model_state_dict']\nmodel_264.load_state_dict(weights_264)\nmodel_264.eval()\n\nmodel_266 = model_238_code.Net(cfg_238).to(cfg_238.device)\nmodel_266 = torch.compile(model_266)\nweights_path_266 = \"/kaggle/input/mabe-weights/MAB-266_best.pth\"\nweights_266 = torch.load(weights_path_266, map_location=cfg_238.device, weights_only=False)#['model_state_dict']\nmodel_266.load_state_dict(weights_266)\nmodel_266.eval()\n\nmodel_267 = model_238_code.Net(cfg_238).to(cfg_238.device)\nmodel_267 = torch.compile(model_267)\nweights_path_267 = \"/kaggle/input/mabe-weights/MAB-267_best.pth\"\nweights_267 = torch.load(weights_path_267, map_location=cfg_238.device, weights_only=False)#['model_state_dict']\nmodel_267.load_state_dict(weights_267)\nmodel_267.eval()\n\nmodel_269 = model_238_code.Net(cfg_238).to(cfg_238.device)\nmodel_269 = torch.compile(model_269)\nweights_path_269 = \"/kaggle/input/mabe-weights/MAB-269_best.pth\"\nweights_269 = torch.load(weights_path_269, map_location=cfg_238.device, weights_only=False)#['model_state_dict']\nmodel_269.load_state_dict(weights_269)\nmodel_269.eval()\n\ncfg_244 = copy(cfg_238)\ncfg_244.per_mouse_feature_dim = 304 // 4\nmodel_244 = model_244_code.Net(cfg_244).to(cfg_244.device)\nmodel_244 = torch.compile(model_244)\nweights_path_244 = \"/kaggle/input/mabe-weights/MAB-244_best.pth\"\nweights_244 = torch.load(weights_path_244, map_location=cfg_244.device, weights_only=False)#['model_state_dict']\nmodel_244.load_state_dict(weights_244)\nmodel_244.eval()\n\ncfg_234 = copy(cfg)\ncfg_234.reverse_time = False\ncfg_234.use_gnn = False\ncfg_234.use_bn = True\ncfg_234.cnn_extractor = True\ncfg_234.encoder_config.input_dim=256\ncfg_234.encoder_config.encoder_dim=256\ncfg_234.encoder_config.num_layers=4\ncfg_234.encoder_config.num_attention_heads=4\ncfg_234.per_mouse_feature_dim = cfg_234.feature_dim // 4\n\nmodel_234 = model_inter.NetInter(cfg_234).to(cfg_234.device)\nmodel_234 = torch.compile(model_234)\nweights_path_234 = \"/kaggle/input/mabe-weights/MAB-234_best.pth\"\nweights_234 = torch.load(weights_path_234, map_location=cfg_234.device, weights_only=False)#['model_state_dict']\nmodel_234.load_state_dict(weights_234)\nmodel_234.eval()\n\ncfg_240 = copy(cfg)\ncfg_240.cnn_extractor = True\ncfg_240.use_gnn = False\ncfg_240.use_bn = True\ncfg_240.encoder_config.input_dim=256\ncfg_240.encoder_config.encoder_dim=256\ncfg_240.encoder_config.num_layers=8\ncfg_240.encoder_config.num_attention_heads=4\ncfg_240.per_mouse_feature_dim = cfg_240.feature_dim // 4\n\nmodel_240 = model_inter.NetInter(cfg_240).to(cfg_240.device)\nmodel_240 = torch.compile(model_240)\nweights_path_240 = \"/kaggle/input/mabe-weights/MAB-240_best.pth\"\nweights_240 = torch.load(weights_path_240, map_location=cfg_240.device, weights_only=False)#['model_state_dict']\nmodel_240.load_state_dict(weights_240)\nmodel_240.eval()\n\ncfg_241 = copy(cfg)\ncfg_241.cnn_extractor = True\ncfg_241.use_gnn = False\ncfg_241.use_bn = True\ncfg_241.encoder_config.input_dim=256\ncfg_241.encoder_config.encoder_dim=256\ncfg_241.encoder_config.num_layers=4\ncfg_241.encoder_config.num_attention_heads=4\ncfg_241.per_mouse_feature_dim = cfg_241.feature_dim // 4\n\nmodel_241 = model_inter.NetInter(cfg_241).to(cfg_241.device)\nmodel_241 = torch.compile(model_241)\nweights_path_241 = \"/kaggle/input/mabe-weights/MAB-241_best.pth\"\nweights_241 = torch.load(weights_path_241, map_location=cfg_241.device, weights_only=False)#['model_state_dict']\nmodel_241.load_state_dict(weights_241)\nmodel_241.eval()\n\ncfg_242 = copy(cfg)\ncfg_242.cnn_extractor = True\ncfg_242.use_gnn = True\ncfg_242.use_bn = True\ncfg_242.encoder_config.input_dim=256\ncfg_242.encoder_config.encoder_dim=256\ncfg_242.encoder_config.num_layers=4\ncfg_242.encoder_config.num_attention_heads=4\ncfg_242.per_mouse_feature_dim = 40 #cfg_242.feature_dim // 4\n\nmodel_242 = model_242_code.Net(cfg_242).to(cfg_242.device)\nmodel_242 = torch.compile(model_242)\nweights_path_242 = \"/kaggle/input/mabe-weights/MAB-242_best.pth\"\nweights_242 = torch.load(weights_path_242, map_location=cfg_242.device, weights_only=False)#['model_state_dict']\nmodel_242.load_state_dict(weights_242)\nmodel_242.eval()\n\ncfg_243 = copy(cfg)\ncfg_243.cnn_extractor = True\ncfg_243.use_gnn = True\ncfg_243.use_bn = True\ncfg_243.encoder_config.input_dim=256\ncfg_243.encoder_config.encoder_dim=256\ncfg_243.encoder_config.num_layers=4\ncfg_243.encoder_config.num_attention_heads=4\ncfg_243.per_mouse_feature_dim = 40 #cfg.feature_dim // len(cfg.set_mice)\n\nmodel_243 = model_243_code.Net(cfg_243).to(cfg_243.device)\nmodel_243 = torch.compile(model_243)\nweights_path_243 = \"/kaggle/input/mabe-weights/MAB-243_best.pth\"\nweights_243 = torch.load(weights_path_243, map_location=cfg_243.device, weights_only=False)#['model_state_dict']\nmodel_243.load_state_dict(weights_243)\nmodel_243.eval()\n\nmodel_294 = model_243_code.Net(cfg_243).to(cfg_243.device)\nmodel_294 = torch.compile(model_294)\nweights_path_294 = \"/kaggle/input/mabe-weights/MAB-294_best.pth\"\nweights_294 = torch.load(weights_path_294, map_location=cfg_243.device, weights_only=False)#['model_state_dict']\nmodel_294.load_state_dict(weights_294)\nmodel_294.eval()\n\nmodel_295 = model_243_code.Net(cfg_243).to(cfg_243.device)\nmodel_295 = torch.compile(model_295)\nweights_path_295 = \"/kaggle/input/mabe-weights/MAB-295_best.pth\"\nweights_295 = torch.load(weights_path_295, map_location=cfg_243.device, weights_only=False)#['model_state_dict']\nmodel_295.load_state_dict(weights_295)\nmodel_295.eval()\n\nmodel_297 = model_243_code.Net(cfg_243).to(cfg_243.device)\nmodel_297 = torch.compile(model_297)\nweights_path_297 = \"/kaggle/input/mabe-weights/MAB-297_best.pth\"\nweights_297 = torch.load(weights_path_297, map_location=cfg_243.device, weights_only=False)#['model_state_dict']\nmodel_297.load_state_dict(weights_297)\nmodel_297.eval()\n\n\ncfg_256 = copy(cfg)\ncfg_256.cnn_extractor = True\ncfg_256.use_gnn = True\ncfg_256.use_bn = True\ncfg_256.encoder_config.input_dim=256\ncfg_256.encoder_config.encoder_dim=256\ncfg_256.encoder_config.num_layers=8\ncfg_256.encoder_config.num_attention_heads=4\n\nmodel_256 = model_inter.NetInter(cfg_256).to(cfg_256.device)\nmodel_256 = torch.compile(model_256)\nweights_path_256 = \"/kaggle/input/mabe-weights/MAB-256_best.pth\"\nweights_256 = torch.load(weights_path_256, map_location=cfg_256.device, weights_only=False)#['model_state_dict']\nmodel_256.load_state_dict(weights_256)\nmodel_256.eval()\n\ncfg_273 = copy(cfg)\ncfg_273.cnn_extractor = True\ncfg_273.use_gnn = False\ncfg_273.use_bn = True\ncfg_273.encoder_config.input_dim=256\ncfg_273.encoder_config.encoder_dim=256\ncfg_273.encoder_config.num_layers=4\ncfg_273.encoder_config.num_attention_heads=4\ncfg_273.per_mouse_feature_dim = cfg_273.feature_dim // 4\n\nmodel_273 = model_gnn.NetGNN(cfg_273).to(cfg_273.device)\nmodel_273 = torch.compile(model_273)\nweights_path_273 = \"/kaggle/input/mabe-weights/MAB-273_best.pth\"\nweights_273 = torch.load(weights_path_273, map_location=cfg_273.device, weights_only=False)#['model_state_dict']\nmodel_273.load_state_dict(weights_273)\nmodel_273.eval() \n\nmodel_274 = model_gnn.NetGNN(cfg_273).to(cfg_273.device)\nmodel_274 = torch.compile(model_274)\nweights_path_274 = \"/kaggle/input/mabe-weights/MAB-274_best.pth\"\nweights_274 = torch.load(weights_path_274, map_location=cfg_273.device, weights_only=False)#['model_state_dict']\nmodel_274.load_state_dict(weights_274)\nmodel_274.eval()\n\nmodel_275 = model_gnn.NetGNN(cfg_273).to(cfg_273.device)\nmodel_275 = torch.compile(model_275)\nweights_path_275 = \"/kaggle/input/mabe-weights/MAB-275_best.pth\"\nweights_275 = torch.load(weights_path_275, map_location=cfg_273.device, weights_only=False)#['model_state_dict']\nmodel_275.load_state_dict(weights_275)\nmodel_275.eval()\n\nmodel_276 = model_gnn.NetGNN(cfg_273).to(cfg_273.device)\nmodel_276 = torch.compile(model_276)\nweights_path_276 = \"/kaggle/input/mabe-weights/MAB-276_best.pth\"\nweights_276 = torch.load(weights_path_276, map_location=cfg_273.device, weights_only=False)#['model_state_dict']\nmodel_276.load_state_dict(weights_276)\nmodel_276.eval()\n\n\ncfg_278 = copy(cfg)\ncfg_278.cnn_extractor = True\ncfg_278.use_gnn = False\ncfg_278.use_bn = True\ncfg_278.encoder_config.input_dim=256\ncfg_278.encoder_config.encoder_dim=256\ncfg_278.encoder_config.num_layers=4\ncfg_278.encoder_config.num_attention_heads=4\ncfg_278.per_mouse_feature_dim = cfg_278.feature_dim // 4\n\nmodel_278 = model_inter.NetInter(cfg_278).to(cfg_278.device)\nmodel_278 = torch.compile(model_278)\nweights_path_278 = \"/kaggle/input/mabe-weights/MAB-278_best.pth\"\nweights_278 = torch.load(weights_path_278, map_location=cfg_278.device, weights_only=False)#['model_state_dict']\nmodel_278.load_state_dict(weights_278)\nmodel_278.eval()\n\nmodel_279 = model_inter.NetInter(cfg_278).to(cfg_278.device)\nmodel_279 = torch.compile(model_279)\nweights_path_279 = \"/kaggle/input/mabe-weights/MAB-279_best.pth\"\nweights_279 = torch.load(weights_path_279, map_location=cfg_278.device, weights_only=False)#['model_state_dict']\nmodel_279.load_state_dict(weights_279)\nmodel_279.eval()\n\nmodel_284 = model_inter.NetInter(cfg_278).to(cfg_278.device)\nmodel_284 = torch.compile(model_284)\nweights_path_284 = \"/kaggle/input/mabe-weights/MAB-284_best.pth\"\nweights_284 = torch.load(weights_path_284, map_location=cfg_278.device, weights_only=False)#['model_state_dict']\nmodel_284.load_state_dict(weights_284)\nmodel_284.eval()\n\nmodel_285 = model_inter.NetInter(cfg_278).to(cfg_278.device)\nmodel_285 = torch.compile(model_285)\nweights_path_285 = \"/kaggle/input/mabe-weights/MAB-285_best.pth\"\nweights_285 = torch.load(weights_path_285, map_location=cfg_278.device, weights_only=False)#['model_state_dict']\nmodel_285.load_state_dict(weights_285)\nmodel_285.eval()\n\n\n# Get action/mouse info from config\nnum_mice = len(cfg.set_mice)\nnum_pairs = num_mice * num_mice\nactions = cfg.set_behavior_classes\nnum_actions = len(actions)\naction_map = cfg.id_to_action_map\n\nmouse_map_str = cfg.mouse_id_to_string\nint_to_str_map = {1: 'mouse1', 2: 'mouse2', 3: 'mouse3', 4: 'mouse4'}\n\nscaler = GradScaler(device=cfg.device) if cfg.mixed_precision else None\ntest_df = pd.read_csv('/kaggle/input/MABe-mouse-behavior-detection/test.csv')\n\nchunk_size = 10  # Adjust based on memory limits; e.g., process 50 videos at a time\nall_predictions_dfs = []\n\nfor i in range(0, len(test_df), chunk_size):\n    chunk_df = test_df.iloc[i:i + chunk_size]\n    temp_prediction_dfs = []\n    val_dataset = CustomDataset(chunk_df, cfg, aug=None, mode=\"val\")\n\n    val_dataloader = DataLoader(\n        val_dataset,\n        shuffle=False,\n        batch_size=cfg.batch_size_val if hasattr(cfg, 'batch_size_val') else cfg.batch_size,\n        num_workers=cfg.num_workers,\n        pin_memory=cfg.pin_memory,\n        collate_fn=val_collate_fn if val_collate_fn else tr_collate_fn,\n        drop_last=False\n    )\n\n    # --- PHASE 1: GPU-BOUND WORK (Inference) ---\n    video_accum = {}  # video_id -> {'all_preds': np.array, 'counts': np.array, 'lab_id': str, 'num_frames': int}\n    for batch in tqdm(val_dataloader, desc=\"Validation\"):\n        batch = batch_to_device(batch, cfg.device)\n        with torch.no_grad():\n            # output_233 = model_233(batch)\n            # preds_batch_233 = output_233['predictions'].cpu().numpy()\n    \n            # output_234 = model_234(batch)\n            # preds_batch_234 = output_234['predictions'].cpu().numpy() \n    \n            # output_238 = model_238(batch)\n            # preds_batch_238 = output_238['predictions'].cpu().numpy()  \n    \n            output_240 = model_240(batch)\n            preds_batch_240 = output_240['predictions'].cpu().numpy()\n        \n            output_242 = model_242(batch)\n            preds_batch_242 = output_242['predictions'].cpu().numpy()\n                \n            # output_243 = model_243(batch)\n            # preds_batch_243 = output_243['predictions'].cpu().numpy()\n                \n            output_244 = model_244(batch)\n            preds_batch_244 = output_244['predictions'].cpu().numpy() \n                \n            # output_245 = model_245(batch)\n            # preds_batch_245 = output_245['predictions'].cpu().numpy() \n                 \n            output_256 = model_256(batch)\n            preds_batch_256 = output_256['predictions'].cpu().numpy()  \n            \n            # ##################\n            output_264 = model_264(batch)\n            preds_batch_264 = output_264['predictions'].cpu().numpy()\n\n            output_266 = model_266(batch)\n            preds_batch_266 = output_266['predictions'].cpu().numpy()\n\n            output_267 = model_267(batch)\n            preds_batch_267 = output_267['predictions'].cpu().numpy()\n\n            output_269 = model_269(batch)\n            preds_batch_269 = output_269['predictions'].cpu().numpy()\n\n            preds_batch_245fold = (preds_batch_264 + preds_batch_266 \n                                + preds_batch_267 + preds_batch_269) / 4.0    \n            ###############\n            output_273 = model_273(batch)\n            preds_batch_273 = output_273['predictions'].cpu().numpy()\n\n            output_274 = model_274(batch)\n            preds_batch_274 = output_274['predictions'].cpu().numpy()\n\n            output_275 = model_275(batch)\n            preds_batch_275 = output_275['predictions'].cpu().numpy()\n\n            output_276 = model_276(batch)\n            preds_batch_276 = output_276['predictions'].cpu().numpy()\n\n            preds_batch_233fold = (preds_batch_273 + preds_batch_274 \n                                + preds_batch_275 + preds_batch_276) / 4.0            \n            ###############\n            output_278 = model_278(batch)\n            preds_batch_278 = output_278['predictions'].cpu().numpy()\n\n            output_279 = model_279(batch)\n            preds_batch_279 = output_279['predictions'].cpu().numpy()\n\n            output_284 = model_284(batch)\n            preds_batch_284 = output_284['predictions'].cpu().numpy()\n\n            output_285 = model_285(batch)\n            preds_batch_285 = output_285['predictions'].cpu().numpy()\n            \n            preds_batch_234fold = (preds_batch_278 + preds_batch_279 \n                                + preds_batch_284 + preds_batch_285) / 4.0\n            #################\n            output_289 = model_289(batch)\n            preds_batch_289 = output_289['predictions'].cpu().numpy()\n            \n            output_290 = model_290(batch)\n            preds_batch_290 = output_290['predictions'].cpu().numpy()\n\n            output_293 = model_293(batch)\n            preds_batch_293 = output_293['predictions'].cpu().numpy()\n\n            output_292 = model_292(batch)\n            preds_batch_292 = output_292['predictions'].cpu().numpy()\n            \n            preds_batch_238fold = (preds_batch_289 + preds_batch_290 \n                                   + preds_batch_293 + preds_batch_292) / 4.0\n            ##################\n            output_294 = model_294(batch)\n            preds_batch_294 = output_294['predictions'].cpu().numpy()\n            \n            output_295 = model_295(batch)\n            preds_batch_295 = output_295['predictions'].cpu().numpy()\n            \n            output_297 = model_297(batch)\n            preds_batch_297 = output_297['predictions'].cpu().numpy()\n            \n            \n            preds_batch_243fold = (preds_batch_294 + preds_batch_295 \n                                   + preds_batch_297) / 3.0\n\n            \n            preds_batch = (preds_batch_233fold + preds_batch_234fold\n                + preds_batch_238fold + preds_batch_240 + preds_batch_242 + preds_batch_243fold\n                + preds_batch_244 + preds_batch_245fold + preds_batch_256) / 9.0\n        \n\n        \n        bs = preds_batch.shape[0]\n        for i in range(bs):\n            video_id = batch['video_id'][i]\n            if video_id not in video_accum:\n                num_frames = batch['num_frames'][i].item()\n                lab_id = val_dataset.video_info[video_id]['lab_id']\n                video_accum[video_id] = {\n                    'all_preds': np.zeros((num_frames, num_pairs, num_actions), dtype=np.float32),\n                    'counts': np.zeros((num_frames, 1, 1), dtype=np.float32),\n                    'lab_id': lab_id,\n                    'num_frames': num_frames\n                }\n            \n            start = batch['start_frame'][i].item()\n            mask = batch['input_mask'][i].cpu().numpy()  # [seq_len]\n            actual_len = int(np.sum(mask))\n            preds = preds_batch[i][:actual_len]\n            # No global cfg.reverse_time check needed; flipping handled above\n            \n            end = min(start + actual_len, video_accum[video_id]['num_frames'])\n            video_accum[video_id]['all_preds'][start:end] += preds[: (end - start)]\n            video_accum[video_id]['counts'][start:end] += 1\n\n    # Average predictions\n    for video_id in video_accum:\n        acc = video_accum[video_id]\n        counts = np.maximum(acc['counts'], 1)\n        acc['all_preds'] /= counts\n\n    # Prepare data for post-processing\n    all_video_data_for_processing = [\n        {\n            'video_preds': video_accum[video_id]['all_preds'],\n            'video_id': video_id,\n            'lab_id': video_accum[video_id]['lab_id'],\n            'num_frames': video_accum[video_id]['num_frames']\n        }\n        for video_id in video_accum\n    ]\n\n    # --- PHASE 2: Post-Processing ---\n    results = Parallel(n_jobs=-1)(\n        delayed(process_one_video_preds_filters)(\n            cfg,\n            data['video_preds'],\n            data['video_id'],\n            data['lab_id'],\n            data['num_frames'],\n            mouse_map_str,\n            action_map,\n            num_pairs,\n            num_actions,\n            min_duration,\n            smooth_sigma=0.0, \n            conf_thresh=0.0,\n            max_gap=0, \n            use_median_filter=False\n        ) for data in tqdm(all_video_data_for_processing, desc=\"Phase 2: Post-Processing\")\n    )\n    temp_prediction_dfs = [df for df in results if df is not None]\n    all_predictions_dfs.extend(temp_prediction_dfs)\n    del val_dataloader\n    del val_dataset\n    gc.collect()\nFINAL_SUBMISSION_COLS = [\n    'row_id',\n    'video_id',\n    'agent_id',\n    'target_id',\n    'action',\n    'start_frame',\n    'stop_frame'\n]\n##END\nif all_predictions_dfs:\n    predictions_df = pd.concat(all_predictions_dfs, ignore_index=True)\n    predictions_df['row_id'] = range(len(predictions_df))\n    predictions_df = predictions_df[FINAL_SUBMISSION_COLS]\nelse:\n    predictions_df = pd.DataFrame(columns=FINAL_SUBMISSION_COLS)\n\n\npredictions_df.to_csv(\"submission.csv\", index=False)\n\nsubmission_file = pd.read_csv('submission.csv')\nprint(f\"Saved the submission file with {len(submission_file)} rows.\")\nend_time = time()\nprint(f\"Total inference and post-processing time: {end_time - start_time:.2f} seconds\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-18T13:06:49.236961Z","iopub.execute_input":"2025-12-18T13:06:49.237272Z","iopub.status.idle":"2025-12-18T13:10:50.916016Z","shell.execute_reply.started":"2025-12-18T13:06:49.237247Z","shell.execute_reply":"2025-12-18T13:10:50.914985Z"}},"outputs":[],"execution_count":null}]}