{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.11.13"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":59156,"databundleVersionId":13874099,"sourceType":"competition"},{"sourceId":13637144,"sourceType":"datasetVersion","datasetId":8662279},{"sourceId":13719557,"sourceType":"datasetVersion","datasetId":8723050},{"sourceId":13855258,"sourceType":"datasetVersion","datasetId":8753460},{"sourceId":13855262,"sourceType":"datasetVersion","datasetId":8826220},{"sourceId":13934704,"sourceType":"datasetVersion","datasetId":8880360},{"sourceId":14111957,"sourceType":"datasetVersion","datasetId":8989282},{"sourceId":14112088,"sourceType":"datasetVersion","datasetId":8989363},{"sourceId":14117334,"sourceType":"datasetVersion","datasetId":8993432},{"sourceId":14164944,"sourceType":"datasetVersion","datasetId":9028829},{"sourceId":14174637,"sourceType":"datasetVersion","datasetId":9035317},{"sourceId":262477103,"sourceType":"kernelVersion"},{"sourceId":272706666,"sourceType":"kernelVersion"}],"dockerImageVersionId":31154,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":374.813175,"end_time":"2025-12-15T14:00:01.496509","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2025-12-15T13:53:46.683334","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# 7th Place Gold - Mouse Comp - [CV 0.559, LB 0.562] - CNN Transformer!\nI am happy to share my 7th place Gold Mouse Comp solution. It is an ensemble of 10x NN. All the NN are CNN Transformers. The ensemble of NN by themselves achieve `CV = 0.546`, `Public LB = 0.551` and `Private LB = 0.518` (rank 8th). Then we ensemble with an XGB stacked over NN and improve to `CV = 0.559`, `Public LB = 0.562` and `Private LB = 0.523` (rank 7th). All 11 inference scripts are in Kaggle dataset [here][2] and all 11 train scripts are in Kaggle dataset folder `train_scripts` [here][3]. Discussion about this solution is [here][1]\n\n![](https://raw.githubusercontent.com/cdeotte/Kaggle_Images/refs/heads/main/Dec-2025/cnn_trans_model.png)\n\n[1]: https://www.kaggle.com/competitions/MABe-mouse-behavior-detection/writeups/7th-place-gold-cnn-transformer-with-invariant-fe\n[2]: https://www.kaggle.com/datasets/cdeotte/mouse-comp-inference-scripts\n[3]: https://www.kaggle.com/datasets/cdeotte/mouse-comp-inference-scripts?select=train_scripts","metadata":{}},{"cell_type":"markdown","source":"# Infer 10 NNs\nAll 10x NN inference scripts can be found in Kaggle dataset [here][1]. And all their train scripts are in Kaggle dataset folder `train_scripts` [here][2]\n\n[1]: https://www.kaggle.com/datasets/cdeotte/mouse-comp-inference-scripts\n[2]: https://www.kaggle.com/datasets/cdeotte/mouse-comp-inference-scripts?select=train_scripts","metadata":{}},{"cell_type":"code","source":"!python /kaggle/input/mouse-comp-inference-scripts/infer-v258x8t.py","metadata":{"_kg_hide-output":true,"collapsed":true,"execution":{"iopub.execute_input":"2025-12-15T13:53:50.53954Z","iopub.status.busy":"2025-12-15T13:53:50.539223Z","iopub.status.idle":"2025-12-15T13:54:40.051923Z","shell.execute_reply":"2025-12-15T13:54:40.050723Z"},"jupyter":{"outputs_hidden":true},"papermill":{"duration":49.522819,"end_time":"2025-12-15T13:54:40.05396","exception":false,"start_time":"2025-12-15T13:53:50.531141","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python /kaggle/input/mouse-comp-inference-scripts/infer-v211dx6t.py","metadata":{"_kg_hide-output":true,"collapsed":true,"execution":{"iopub.execute_input":"2025-12-15T13:54:40.076572Z","iopub.status.busy":"2025-12-15T13:54:40.076304Z","iopub.status.idle":"2025-12-15T13:55:12.652779Z","shell.execute_reply":"2025-12-15T13:55:12.651858Z"},"jupyter":{"outputs_hidden":true},"papermill":{"duration":32.587595,"end_time":"2025-12-15T13:55:12.654231","exception":false,"start_time":"2025-12-15T13:54:40.066636","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!rm -rf test_tracking3","metadata":{"execution":{"iopub.execute_input":"2025-12-15T13:55:12.673339Z","iopub.status.busy":"2025-12-15T13:55:12.673074Z","iopub.status.idle":"2025-12-15T13:55:12.806214Z","shell.execute_reply":"2025-12-15T13:55:12.805124Z"},"papermill":{"duration":0.144029,"end_time":"2025-12-15T13:55:12.807619","exception":false,"start_time":"2025-12-15T13:55:12.66359","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python /kaggle/input/mouse-comp-inference-scripts/infer-v315x6t.py","metadata":{"_kg_hide-output":true,"collapsed":true,"execution":{"iopub.execute_input":"2025-12-15T13:55:12.826692Z","iopub.status.busy":"2025-12-15T13:55:12.826029Z","iopub.status.idle":"2025-12-15T13:55:31.527801Z","shell.execute_reply":"2025-12-15T13:55:31.526966Z"},"jupyter":{"outputs_hidden":true},"papermill":{"duration":18.712791,"end_time":"2025-12-15T13:55:31.52923","exception":false,"start_time":"2025-12-15T13:55:12.816439","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python /kaggle/input/mouse-comp-inference-scripts/infer-v330x6t.py","metadata":{"_kg_hide-output":true,"collapsed":true,"execution":{"iopub.execute_input":"2025-12-15T13:55:31.549939Z","iopub.status.busy":"2025-12-15T13:55:31.549314Z","iopub.status.idle":"2025-12-15T13:55:56.114696Z","shell.execute_reply":"2025-12-15T13:55:56.113904Z"},"jupyter":{"outputs_hidden":true},"papermill":{"duration":24.576908,"end_time":"2025-12-15T13:55:56.116102","exception":false,"start_time":"2025-12-15T13:55:31.539194","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python /kaggle/input/mouse-comp-inference-scripts/infer-v357x6t.py","metadata":{"_kg_hide-output":true,"collapsed":true,"execution":{"iopub.execute_input":"2025-12-15T13:55:56.137514Z","iopub.status.busy":"2025-12-15T13:55:56.137052Z","iopub.status.idle":"2025-12-15T13:56:21.517831Z","shell.execute_reply":"2025-12-15T13:56:21.517089Z"},"jupyter":{"outputs_hidden":true},"papermill":{"duration":25.392796,"end_time":"2025-12-15T13:56:21.519292","exception":false,"start_time":"2025-12-15T13:55:56.126496","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python /kaggle/input/mouse-comp-inference-scripts/infer-v340x6t.py","metadata":{"_kg_hide-output":true,"collapsed":true,"execution":{"iopub.execute_input":"2025-12-15T13:56:21.542609Z","iopub.status.busy":"2025-12-15T13:56:21.542038Z","iopub.status.idle":"2025-12-15T13:56:47.205382Z","shell.execute_reply":"2025-12-15T13:56:47.20464Z"},"jupyter":{"outputs_hidden":true},"papermill":{"duration":25.676143,"end_time":"2025-12-15T13:56:47.206752","exception":false,"start_time":"2025-12-15T13:56:21.530609","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python /kaggle/input/mouse-comp-inference-scripts/infer-v345x6t.py","metadata":{"_kg_hide-output":true,"collapsed":true,"execution":{"iopub.execute_input":"2025-12-15T13:56:47.231876Z","iopub.status.busy":"2025-12-15T13:56:47.231634Z","iopub.status.idle":"2025-12-15T13:57:14.078353Z","shell.execute_reply":"2025-12-15T13:57:14.077308Z"},"jupyter":{"outputs_hidden":true},"papermill":{"duration":26.860676,"end_time":"2025-12-15T13:57:14.079849","exception":false,"start_time":"2025-12-15T13:56:47.219173","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python /kaggle/input/mouse-comp-inference-scripts/infer-v417x6t.py","metadata":{"_kg_hide-output":true,"collapsed":true,"execution":{"iopub.execute_input":"2025-12-15T13:57:14.105858Z","iopub.status.busy":"2025-12-15T13:57:14.105621Z","iopub.status.idle":"2025-12-15T13:57:35.632095Z","shell.execute_reply":"2025-12-15T13:57:35.63135Z"},"jupyter":{"outputs_hidden":true},"papermill":{"duration":21.540726,"end_time":"2025-12-15T13:57:35.633421","exception":false,"start_time":"2025-12-15T13:57:14.092695","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python /kaggle/input/mouse-comp-inference-scripts/infer-v421x6t.py","metadata":{"_kg_hide-output":true,"collapsed":true,"execution":{"iopub.execute_input":"2025-12-15T13:57:35.66178Z","iopub.status.busy":"2025-12-15T13:57:35.661194Z","iopub.status.idle":"2025-12-15T13:57:58.743964Z","shell.execute_reply":"2025-12-15T13:57:58.743235Z"},"jupyter":{"outputs_hidden":true},"papermill":{"duration":23.098154,"end_time":"2025-12-15T13:57:58.745278","exception":false,"start_time":"2025-12-15T13:57:35.647124","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python /kaggle/input/mouse-comp-inference-scripts/infer-v422x6t.py","metadata":{"_kg_hide-output":true,"collapsed":true,"execution":{"iopub.execute_input":"2025-12-15T13:57:58.775666Z","iopub.status.busy":"2025-12-15T13:57:58.775388Z","iopub.status.idle":"2025-12-15T13:58:23.273584Z","shell.execute_reply":"2025-12-15T13:58:23.272575Z"},"jupyter":{"outputs_hidden":true},"papermill":{"duration":24.514538,"end_time":"2025-12-15T13:58:23.275035","exception":false,"start_time":"2025-12-15T13:57:58.760497","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!rm submission.csv","metadata":{"execution":{"iopub.execute_input":"2025-12-15T13:58:23.307307Z","iopub.status.busy":"2025-12-15T13:58:23.306598Z","iopub.status.idle":"2025-12-15T13:58:23.42119Z","shell.execute_reply":"2025-12-15T13:58:23.420463Z"},"papermill":{"duration":0.131521,"end_time":"2025-12-15T13:58:23.422269","exception":false,"start_time":"2025-12-15T13:58:23.290748","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TT_PER_LAB_NN = {'AdaptableSnail': {'approach': 0.78,\n  'attack': 0.59,\n  'avoid': 0.63,\n  'chase': 0.61,\n  'chaseattack': 0.6900000000000001,\n  'rear': 0.7000000000000001,\n  'submit': 0.22000000000000008},\n 'BoisterousParrot': {'shepherd': 0.7100000000000001},\n 'CRIM13': {'approach': 0.79,\n  'attack': 0.59,\n  'disengage': 0.7200000000000001,\n  'mount': 0.48000000000000004,\n  'rear': 0.58,\n  'selfgroom': 0.67,\n  'sniff': 0.6900000000000001},\n 'CalMS21_supplemental': {'approach': 0.89,\n  'attack': 0.89,\n  'attemptmount': 0.6,\n  'dominancemount': 0.79,\n  'intromit': 0.7300000000000001,\n  'mount': 0.8300000000000001,\n  'sniff': 0.7300000000000001,\n  'sniffbody': 0.76,\n  'sniffface': 0.88,\n  'sniffgenital': 0.68},\n 'CalMS21_task1': {'approach': 0.9,\n  'attack': 0.88,\n  'genitalgroom': 0.7100000000000001,\n  'intromit': 0.7100000000000001,\n  'mount': 0.8400000000000001,\n  'sniff': 0.7000000000000001,\n  'sniffbody': 0.6900000000000001,\n  'sniffface': 0.87,\n  'sniffgenital': 0.85},\n 'CalMS21_task2': {'attack': 0.87,\n  'mount': 0.12000000000000001,\n  'sniff': 0.63},\n 'CautiousGiraffe': {'chase': 0.3200000000000001,\n  'escape': 0.8,\n  'reciprocalsniff': 0.8500000000000001,\n  'sniff': 0.64,\n  'sniffbody': 0.54,\n  'sniffgenital': 0.7100000000000001},\n 'DeliriousFly': {'attack': 0.61,\n  'dominance': 0.49000000000000005,\n  'sniff': 0.51},\n 'ElegantMink': {'allogroom': 0.35000000000000014,\n  'attack': 0.49000000000000005,\n  'attemptmount': 0.6900000000000001,\n  'ejaculate': 0.18000000000000005,\n  'intromit': 0.57,\n  'mount': 0.62,\n  'sniff': 0.38000000000000006},\n 'GroovyShrew': {'approach': 0.86,\n  'attemptmount': 0.34000000000000014,\n  'climb': 0.30000000000000004,\n  'defend': 0.5,\n  'dig': 0.67,\n  'escape': 0.7100000000000001,\n  'rear': 0.76,\n  'rest': 0.61,\n  'run': 0.7200000000000001,\n  'selfgroom': 0.76,\n  'sniff': 0.62,\n  'sniffgenital': 0.7200000000000001},\n 'InvincibleJellyfish': {'allogroom': 0.35000000000000014,\n  'attack': 0.7200000000000001,\n  'dig': 0.65,\n  'dominancegroom': 0.54,\n  'escape': 0.35000000000000014,\n  'selfgroom': 0.49000000000000005,\n  'sniff': 0.67,\n  'sniffgenital': 0.68},\n 'JovialSwallow': {'attack': 0.7200000000000001,\n  'chase': 0.2800000000000001,\n  'sniff': 0.6900000000000001},\n 'LyricalHare': {'approach': 0.25000000000000006,\n  'attack': 0.77,\n  'defend': 0.62,\n  'escape': 0.8300000000000001,\n  'freeze': 0.2,\n  'rear': 0.7000000000000001,\n  'sniff': 0.61},\n 'NiftyGoldfinch': {'approach': 0.87,\n  'attack': 0.7400000000000001,\n  'biteobject': 0.37000000000000005,\n  'chase': 0.86,\n  'climb': 0.67,\n  'defend': 0.66,\n  'dig': 0.76,\n  'escape': 0.87,\n  'exploreobject': 0.57,\n  'flinch': 0.64,\n  'follow': 0.7300000000000001,\n  'rear': 0.63,\n  'selfgroom': 0.7100000000000001,\n  'sniff': 0.58,\n  'sniffface': 0.88,\n  'sniffgenital': 0.4100000000000001,\n  'tussle': 0.45},\n 'PleasantMeerkat': {'attack': 0.59,\n  'chase': 0.8200000000000001,\n  'escape': 0.6900000000000001,\n  'follow': 0.67},\n 'ReflectiveManatee': {'attack': 0.86, 'sniff': 0.68},\n 'SparklingTapir': {'attack': 0.86,\n  'defend': 0.88,\n  'escape': 0.9,\n  'mount': 0.8300000000000001},\n 'TranquilPanther': {'intromit': 0.63,\n  'mount': 0.7500000000000001,\n  'rear': 0.5599999999999999,\n  'selfgroom': 0.5,\n  'sniff': 0.62,\n  'sniffgenital': 0.65},\n 'UppityFerret': {'huddle': 0.65,\n  'reciprocalsniff': 0.81,\n  'sniffgenital': 0.7200000000000001},\n 'unknown': {'allogroom': 0.34000000000000014, 'approach': 0.9400000000000001, 'attack': 0.77, 'attemptmount': 0.53, 'avoid': 0.75, 'biteobject': 0.55, 'chase': 0.8400000000000001, 'chaseattack': 0.76, 'climb': 0.48000000000000004, 'defend': 0.59, 'dig': 0.8, 'disengage': 0.8400000000000001, 'dominance': 0.63, 'dominancegroom': 0.66, 'dominancemount': 0.93, 'ejaculate': 0.35000000000000003, 'escape': 0.8300000000000001, 'exploreobject': 0.7400000000000001, 'flinch': 0.78, 'follow': 0.89, 'freeze': 0.30000000000000004, 'genitalgroom': 0.59, 'huddle': 0.6, 'intromit': 0.49000000000000005, 'mount': 0.7500000000000001, 'rear': 0.64, 'reciprocalsniff': 0.87, 'rest': 0.66, 'run': 0.89, 'selfgroom': 0.8, 'shepherd': 0.88, 'sniff': 0.61, 'sniffbody': 0.63, 'sniffface': 0.91, 'sniffgenital': 0.68, 'submit': 0.37000000000000005, 'tussle': 0.5499999999999999}}","metadata":{"_kg_hide-input":false,"execution":{"iopub.execute_input":"2025-12-15T13:58:23.790125Z","iopub.status.busy":"2025-12-15T13:58:23.789827Z","iopub.status.idle":"2025-12-15T13:58:23.800713Z","shell.execute_reply":"2025-12-15T13:58:23.800184Z"},"papermill":{"duration":0.028857,"end_time":"2025-12-15T13:58:23.801802","exception":false,"start_time":"2025-12-15T13:58:23.772945","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TRAIN_LAB_ACTIONS = {\n    \"AdaptableSnail\": [\n        \"approach\", \"attack\", \"avoid\", \"chase\", \"chaseattack\", \"rear\", \"submit\"\n    ],\n    \"BoisterousParrot\": [\n        \"shepherd\"\n    ],\n    \"CRIM13\": [\n        \"approach\", \"attack\", \"disengage\", \"mount\", \"rear\",\n        \"selfgroom\", \"sniff\"\n    ],\n    \"CalMS21_supplemental\": [\n        \"approach\", \"attemptmount\", \"attack\", \"dominancemount\", \"intromit\",\n        \"mount\", \"sniff\", \"sniffbody\", \"sniffgenital\", \"sniffface\"\n    ],\n    \"CalMS21_task1\": [\n        \"approach\", \"attack\", \"genitalgroom\", \"intromit\", \"mount\",\n        \"sniff\", \"sniffbody\", \"sniffgenital\", \"sniffface\"\n    ],\n    \"CalMS21_task2\": [\n        \"attack\", \"mount\", \"sniff\"\n    ],\n    \"CautiousGiraffe\": [\n        \"chase\", \"escape\", \"reciprocalsniff\", \"sniff\", \"sniffbody\", \"sniffgenital\"\n    ],\n    \"DeliriousFly\": [\n        \"attack\", \"dominance\", \"sniff\"\n    ],\n    \"ElegantMink\": [\n        \"allogroom\", \"attack\", \"attemptmount\", \"ejaculate\",\n        \"intromit\", \"mount\", \"sniff\"\n    ],\n    \"GroovyShrew\": [\n        \"attemptmount\", \"climb\", \"defend\", \"dig\", \"escape\", \"rear\",\n        \"rest\", \"run\", \"selfgroom\", \"sniff\", \"sniffgenital\", \"approach\"\n    ],\n    \"InvincibleJellyfish\": [\n        \"allogroom\", \"attack\", \"dig\", \"dominancegroom\", \"escape\",\n        \"selfgroom\", \"sniff\", \"sniffgenital\"\n    ],\n    \"JovialSwallow\": [\n        \"attack\", \"chase\", \"sniff\"\n    ],\n    \"LyricalHare\": [\n        \"approach\", \"attack\", \"defend\", \"escape\", \"freeze\", \"rear\", \"sniff\"\n    ],\n    \"NiftyGoldfinch\": [\n        \"approach\", \"attack\", \"biteobject\", \"chase\", \"climb\", \"defend\",\n        \"dig\", \"escape\", \"exploreobject\", \"flinch\", \"follow\", \"rear\",\n        \"run\", \"selfgroom\", \"sniff\", \"sniffgenital\", \"sniffface\", \"tussle\"\n    ],\n    \"PleasantMeerkat\": [\n        \"attack\", \"chase\", \"escape\", \"follow\"\n    ],\n    \"ReflectiveManatee\": [\n        \"attack\", \"sniff\"\n    ],\n    \"SparklingTapir\": [\n        \"attack\", \"defend\", \"escape\", \"mount\"\n    ],\n    \"TranquilPanther\": [\n        \"intromit\", \"mount\", \"rear\", \"selfgroom\", \"sniff\", \"sniffgenital\"\n    ],\n    \"UppityFerret\": [\n        \"huddle\", \"reciprocalsniff\", \"sniffgenital\"\n    ]\n}","metadata":{"execution":{"iopub.execute_input":"2025-12-15T13:58:23.914898Z","iopub.status.busy":"2025-12-15T13:58:23.914721Z","iopub.status.idle":"2025-12-15T13:58:23.920461Z","shell.execute_reply":"2025-12-15T13:58:23.919793Z"},"papermill":{"duration":0.022735,"end_time":"2025-12-15T13:58:23.921547","exception":false,"start_time":"2025-12-15T13:58:23.898812","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"SELF_ACTIONS = ['rear', 'selfgroom', 'genitalgroom', 'rest', 'climb', 'dig', 'run', 'freeze', 'biteobject', 'exploreobject', 'huddle']\nPAIR_ACTIONS = ['approach', 'attack', 'avoid', 'chase', 'chaseattack', 'submit', 'shepherd', 'disengage', 'mount', 'sniff', 'sniffgenital', 'dominancemount', 'sniffbody', 'sniffface', 'attemptmount', 'intromit', 'escape', 'reciprocalsniff', 'dominance', 'allogroom', 'ejaculate', 'defend', 'dominancegroom', 'flinch', 'follow', 'tussle']","metadata":{"execution":{"iopub.execute_input":"2025-12-15T13:58:23.95246Z","iopub.status.busy":"2025-12-15T13:58:23.951959Z","iopub.status.idle":"2025-12-15T13:58:23.955707Z","shell.execute_reply":"2025-12-15T13:58:23.955199Z"},"papermill":{"duration":0.02018,"end_time":"2025-12-15T13:58:23.956762","exception":false,"start_time":"2025-12-15T13:58:23.936582","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TIE_CONFIG_V2 = {'InvincibleJellyfish': {'boost': {'allogroom': 0.07026408914940072,\n   'dominancegroom': 0.04895076820557969},\n  'penalize': {'sniff': 0.032264889275311445},\n  'prefer': [('allogroom', 'sniff', 0.03540325922834654),\n   ('dominancegroom', 'sniff', 0.03191512180892245)]},\n 'CautiousGiraffe': {'boost': {'chase': 0.046873419525772175,\n   'sniffbody': 0.04956469609402715},\n  'penalize': {'sniffgenital': 0.026045976124565565},\n  'prefer': [('chase', 'sniffgenital', 0.04373521192229624),\n   ('sniffbody', 'reciprocalsniff', 0.042224058799378184)]},\n 'ElegantMink': {'boost': {'allogroom': 0.051344438992756466,\n   'ejaculate': 0.04419288246136156},\n  'penalize': {'sniff': 0.028419975022486547},\n  'prefer': [('allogroom', 'sniff', 0.04197073992987219),\n   ('ejaculate', 'intromit', 0.03535669428927513)]},\n 'NiftyGoldfinch': {'boost': {'tussle': 0.051902678013579284,\n   'sniffgenital': 0.03182946515560948,\n   'biteobject': 0.04947620077069683},\n  'penalize': {'rear': 0.021486118867211898,\n   'selfgroom': 0.012126113327424002},\n  'prefer': [('tussle', 'defend', 0.033923677384710575),\n   ('sniffgenital', 'sniff', 0.0441007770763866),\n   ('biteobject', 'sniff', 0.035835949429890365)]},\n 'AdaptableSnail': {'boost': {'chaseattack': 0.04, 'chase': 0.035},\n  'penalize': {'avoid': 0.025},\n  'prefer': [('chaseattack', 'chase', 0.035), ('chase', 'avoid', 0.035)]}}\n\nTIE_CONFIG_V2 = None","metadata":{"execution":{"iopub.execute_input":"2025-12-15T13:58:23.877174Z","iopub.status.busy":"2025-12-15T13:58:23.876994Z","iopub.status.idle":"2025-12-15T13:58:23.882443Z","shell.execute_reply":"2025-12-15T13:58:23.881902Z"},"papermill":{"duration":0.023431,"end_time":"2025-12-15T13:58:23.883555","exception":false,"start_time":"2025-12-15T13:58:23.860124","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Infer XGBoost\nThis XGBoost is stacked over NN-v258, NN-v315, and NN-v421. It uses the features from Mahdi Ravaghi public notebook [here][1]. The inference script is in my inference script Kaggle dataset [here][2]. And the train script is [here][3] (named `xgb-250c.ipynb`)\n\n[1]: https://www.kaggle.com/code/ravaghi/social-action-recognition-in-mice-xgboost\n[2]: https://www.kaggle.com/datasets/cdeotte/mouse-comp-inference-scripts?select=infer-xgb-250c.py\n[3]: https://www.kaggle.com/datasets/cdeotte/mouse-comp-inference-scripts?select=train_scripts","metadata":{"papermill":{"duration":0.014543,"end_time":"2025-12-15T13:58:23.451811","exception":false,"start_time":"2025-12-15T13:58:23.437268","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import pandas as pd\ndf = pd.read_csv(\"/kaggle/input/MABe-mouse-behavior-detection/test.csv\")\nV2L = df[['lab_id','video_id']].drop_duplicates().set_index('video_id').lab_id.to_dict()\n\nXGB_WGT = 0.4","metadata":{"execution":{"iopub.execute_input":"2025-12-15T13:58:23.482984Z","iopub.status.busy":"2025-12-15T13:58:23.482765Z","iopub.status.idle":"2025-12-15T13:58:23.75403Z","shell.execute_reply":"2025-12-15T13:58:23.753366Z"},"papermill":{"duration":0.288678,"end_time":"2025-12-15T13:58:23.755515","exception":false,"start_time":"2025-12-15T13:58:23.466837","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TT_PER_LAB_XGB = {'AdaptableSnail': {'approach': 0.09,\n  'attack': 0.1,\n  'avoid': 0.13,\n  'chase': 0.09,\n  'chaseattack': 0.02,\n  'rear': 0.18999999999999995,\n  'submit': 0.01},\n 'BoisterousParrot': {'shepherd': 0.3200000000000002},\n 'CRIM13': {'approach': 0.3200000000000002,\n  'attack': 0.02,\n  'disengage': 0.3300000000000002,\n  'mount': 0.12,\n  'rear': 0.16999999999999998,\n  'selfgroom': 0.13,\n  'sniff': 0.4},\n 'CalMS21_supplemental': {'approach': 0.28000000000000014,\n  'attack': 0.4600000000000001,\n  'attemptmount': 0.05,\n  'dominancemount': 0.28000000000000014,\n  'intromit': 0.43000000000000016,\n  'mount': 0.4,\n  'sniff': 0.3900000000000001,\n  'sniffbody': 0.3700000000000001,\n  'sniffface': 0.3900000000000001,\n  'sniffgenital': 0.30000000000000004},\n 'CalMS21_task1': {'approach': 0.2,\n  'attack': 0.45000000000000007,\n  'genitalgroom': 0.3900000000000001,\n  'intromit': 0.12,\n  'mount': 0.5400000000000001,\n  'sniff': 0.43000000000000016,\n  'sniffbody': 0.3200000000000002,\n  'sniffface': 0.3300000000000002,\n  'sniffgenital': 0.44000000000000017},\n 'CalMS21_task2': {'attack': 0.4,\n  'mount': 0.44000000000000006,\n  'sniff': 0.45000000000000007},\n 'CautiousGiraffe': {'chase': 0.04,\n  'escape': 0.09,\n  'reciprocalsniff': 0.44000000000000017,\n  'sniff': 0.27000000000000013,\n  'sniffbody': 0.13,\n  'sniffgenital': 0.2},\n 'DeliriousFly': {'attack': 0.2, 'dominance': 0.2, 'sniff': 0.12},\n 'ElegantMink': {'allogroom': 0.13,\n  'attack': 0.3800000000000001,\n  'attemptmount': 0.13999999999999999,\n  'ejaculate': 0.04,\n  'intromit': 0.3300000000000002,\n  'mount': 0.24999999999999992,\n  'sniff': 0.23999999999999994},\n 'GroovyShrew': {'approach': 0.21999999999999995,\n  'attemptmount': 0.01,\n  'climb': 0.14999999999999997,\n  'defend': 0.01,\n  'dig': 0.30000000000000004,\n  'escape': 0.14,\n  'rear': 0.4,\n  'rest': 0.16999999999999998,\n  'run': 0.07,\n  'selfgroom': 0.2500000000000001,\n  'sniff': 0.3700000000000001,\n  'sniffgenital': 0.18999999999999995},\n 'InvincibleJellyfish': {'allogroom': 0.01,\n  'attack': 0.15999999999999998,\n  'dig': 0.06,\n  'dominancegroom': 0.03,\n  'escape': 0.01,\n  'selfgroom': 0.01,\n  'sniff': 0.31000000000000016,\n  'sniffgenital': 0.17999999999999997},\n 'JovialSwallow': {'attack': 0.42000000000000015,\n  'chase': 0.2,\n  'sniff': 0.3500000000000001},\n 'LyricalHare': {'approach': 0.09,\n  'attack': 0.5500000000000002,\n  'defend': 0.22999999999999995,\n  'escape': 0.3800000000000001,\n  'freeze': 0.2600000000000001,\n  'rear': 0.2500000000000001,\n  'sniff': 0.12},\n 'NiftyGoldfinch': {'approach': 0.22999999999999995,\n  'attack': 0.2600000000000001,\n  'biteobject': 0.01,\n  'chase': 0.11,\n  'climb': 0.09,\n  'defend': 0.20999999999999996,\n  'dig': 0.29000000000000015,\n  'escape': 0.3600000000000001,\n  'exploreobject': 0.01,\n  'flinch': 0.04,\n  'follow': 0.11,\n  'rear': 0.1,\n  'selfgroom': 0.29000000000000015,\n  'sniff': 0.2500000000000001,\n  'sniffface': 0.27000000000000013,\n  'sniffgenital': 0.04,\n  'tussle': 0.01},\n 'PleasantMeerkat': {'attack': 0.1,\n  'chase': 0.02,\n  'escape': 0.03,\n  'follow': 0.24000000000000007},\n 'ReflectiveManatee': {'attack': 0.4500000000000002,\n  'sniff': 0.4700000000000001},\n 'SparklingTapir': {'attack': 0.41000000000000014,\n  'defend': 0.22000000000000006,\n  'escape': 0.3500000000000002,\n  'mount': 0.28000000000000014},\n 'TranquilPanther': {'intromit': 0.30000000000000004,\n  'mount': 0.22999999999999995,\n  'rear': 0.15999999999999998,\n  'selfgroom': 0.04,\n  'sniff': 0.20999999999999996,\n  'sniffgenital': 0.23999999999999994},\n 'UppityFerret': {'huddle': 0.20999999999999996,\n  'reciprocalsniff': 0.3400000000000002,\n  'sniffgenital': 0.3400000000000002},\n 'unknown': {'allogroom': 0.09999999999999999, 'approach': 0.20000000000000004, 'attack': 0.18000000000000002, 'attemptmount': 0.08000000000000002, 'avoid': 0.13999999999999996, 'biteobject': 0.03, 'chase': 0.05, 'chaseattack': 0.05, 'climb': 0.12999999999999998, 'defend': 0.21, 'dig': 0.15, 'disengage': 0.25, 'dominance': 0.09, 'dominancegroom': 0.07, 'dominancemount': 0.27, 'ejaculate': 0.09000000000000001, 'escape': 0.20000000000000004, 'exploreobject': 0.06, 'flinch': 0.03, 'follow': 0.20000000000000004, 'freeze': 0.06, 'genitalgroom': 0.06, 'huddle': 0.18, 'intromit': 0.19000000000000003, 'mount': 0.25, 'rear': 0.16, 'reciprocalsniff': 0.33000000000000007, 'rest': 0.20000000000000004, 'run': 0.07, 'selfgroom': 0.14999999999999997, 'shepherd': 0.26, 'sniff': 0.29000000000000004, 'sniffbody': 0.21, 'sniffface': 0.36, 'sniffgenital': 0.26, 'submit': 0.03, 'tussle': 0.03} }","metadata":{"_kg_hide-input":false,"execution":{"iopub.execute_input":"2025-12-15T13:58:23.83333Z","iopub.status.busy":"2025-12-15T13:58:23.833109Z","iopub.status.idle":"2025-12-15T13:58:23.843988Z","shell.execute_reply":"2025-12-15T13:58:23.843377Z"},"papermill":{"duration":0.027991,"end_time":"2025-12-15T13:58:23.845069","exception":false,"start_time":"2025-12-15T13:58:23.817078","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install --no-index --find-links=/kaggle/input/pip-install-koolbox koolbox","metadata":{"_kg_hide-output":true,"collapsed":true,"execution":{"iopub.execute_input":"2025-12-15T13:58:24.017034Z","iopub.status.busy":"2025-12-15T13:58:24.016675Z","iopub.status.idle":"2025-12-15T13:58:31.777875Z","shell.execute_reply":"2025-12-15T13:58:31.777166Z"},"jupyter":{"outputs_hidden":true},"papermill":{"duration":7.777879,"end_time":"2025-12-15T13:58:31.779109","exception":false,"start_time":"2025-12-15T13:58:24.00123","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!python /kaggle/input/mouse-comp-inference-scripts/infer-xgb-250c.py","metadata":{"_kg_hide-output":true,"collapsed":true,"execution":{"iopub.execute_input":"2025-12-15T13:58:31.812963Z","iopub.status.busy":"2025-12-15T13:58:31.812317Z","iopub.status.idle":"2025-12-15T13:59:21.32837Z","shell.execute_reply":"2025-12-15T13:59:21.327444Z"},"jupyter":{"outputs_hidden":true},"papermill":{"duration":49.534247,"end_time":"2025-12-15T13:59:21.330015","exception":false,"start_time":"2025-12-15T13:58:31.795768","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd, numpy as np\ndd = pd.read_parquet(\"/tmp/submission2.pqt\")\ndd2 = pd.read_parquet(\"/tmp/submission3.pqt\")\ndd = pd.concat([dd,dd2],axis=1)\ndel dd2\nprint( dd.shape )\ndd.head()","metadata":{"collapsed":true,"execution":{"iopub.execute_input":"2025-12-15T13:59:21.423891Z","iopub.status.busy":"2025-12-15T13:59:21.423701Z","iopub.status.idle":"2025-12-15T13:59:21.530649Z","shell.execute_reply":"2025-12-15T13:59:21.529809Z"},"jupyter":{"outputs_hidden":true},"papermill":{"duration":0.128364,"end_time":"2025-12-15T13:59:21.531689","exception":false,"start_time":"2025-12-15T13:59:21.403325","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"COLS = ['action_allogroom', 'action_approach', 'action_attack', 'action_attemptmount', 'action_avoid', 'action_biteobject', 'action_chase', 'action_chaseattack', 'action_climb', 'action_defend', 'action_dig', 'action_disengage', 'action_dominance', 'action_dominancegroom', 'action_dominancemount', 'action_ejaculate', 'action_escape', 'action_exploreobject', 'action_flinch', 'action_follow', 'action_freeze', 'action_genitalgroom', 'action_huddle', 'action_intromit', 'action_mount', 'action_rear', 'action_reciprocalsniff', 'action_rest', 'action_run', 'action_selfgroom', 'action_shepherd', 'action_sniff', 'action_sniffbody', 'action_sniffface', 'action_sniffgenital', 'action_submit', 'action_tussle', 'video_id', 'agent_id', 'target_id', 'frame']\nACTIONS = ['allogroom', 'approach', 'attack', 'attemptmount', 'avoid', 'biteobject', 'chase', 'chaseattack', 'climb', 'defend', 'dig', 'disengage', 'dominance', 'dominancegroom', 'dominancemount', 'ejaculate', 'escape', 'exploreobject', 'flinch', 'follow', 'freeze', 'genitalgroom', 'huddle', 'intromit', 'mount', 'rear', 'reciprocalsniff', 'rest', 'run', 'selfgroom', 'shepherd', 'sniff', 'sniffbody', 'sniffface', 'sniffgenital', 'submit', 'tussle']","metadata":{"execution":{"iopub.execute_input":"2025-12-15T13:59:21.573808Z","iopub.status.busy":"2025-12-15T13:59:21.57358Z","iopub.status.idle":"2025-12-15T13:59:21.578348Z","shell.execute_reply":"2025-12-15T13:59:21.577658Z"},"papermill":{"duration":0.02713,"end_time":"2025-12-15T13:59:21.579603","exception":false,"start_time":"2025-12-15T13:59:21.552473","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"ACTS = []\nfor c in dd.columns[:-4]:\n    n = f\"action_{c}\"\n    dd = dd.rename({c:n},axis=1)\n    dd[n] = dd[n].astype('float32')\n    ACTS.append(c)\nfor c in ACTIONS:\n    if c in ACTS: continue\n    n = f\"action_{c}\"\n    dd[n] = np.float32(0)\ndd[\"agent_id\"] = dd.agent_id.str.replace(\"mouse\",\"\").astype('int32')\ndd[\"target_id\"] = dd.target_id.str.replace(\"mouse\",\"\")\ndd.loc[dd.target_id=='self','target_id'] = dd.loc[dd.target_id=='self','agent_id'].astype('str')\ndd[\"target_id\"] = dd[\"target_id\"].astype('int32')\ndd = dd.rename({\"video_frame\":\"frame\"},axis=1)\ndd[\"frame\"] = dd[\"frame\"].astype('int32')\ndd = dd[COLS]\ndd = dd.fillna(0)\nprint( dd.shape )\ndd.head()","metadata":{"execution":{"iopub.execute_input":"2025-12-15T13:59:21.620609Z","iopub.status.busy":"2025-12-15T13:59:21.620382Z","iopub.status.idle":"2025-12-15T13:59:22.112925Z","shell.execute_reply":"2025-12-15T13:59:22.112009Z"},"papermill":{"duration":0.514135,"end_time":"2025-12-15T13:59:22.114136","exception":false,"start_time":"2025-12-15T13:59:21.600001","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!mkdir /tmp/xgb_preds\n!rm /tmp/submission2.pqt /tmp/submission3.pqt","metadata":{"execution":{"iopub.execute_input":"2025-12-15T13:59:22.158034Z","iopub.status.busy":"2025-12-15T13:59:22.157454Z","iopub.status.idle":"2025-12-15T13:59:22.404834Z","shell.execute_reply":"2025-12-15T13:59:22.403814Z"},"papermill":{"duration":0.269777,"end_time":"2025-12-15T13:59:22.406243","exception":false,"start_time":"2025-12-15T13:59:22.136466","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"V = dd.video_id.unique()\nfor v in V:\n    tmp = dd.loc[dd.video_id==v]\n    n = f\"p{v}.pqt\"\n    tmp.to_parquet(f\"/tmp/xgb_preds/{n}\",index=False)\n    print(f\"{v}, \",end=\"\")\nprint()","metadata":{"execution":{"iopub.execute_input":"2025-12-15T13:59:22.448446Z","iopub.status.busy":"2025-12-15T13:59:22.447868Z","iopub.status.idle":"2025-12-15T13:59:22.83167Z","shell.execute_reply":"2025-12-15T13:59:22.830764Z"},"papermill":{"duration":0.406164,"end_time":"2025-12-15T13:59:22.832961","exception":false,"start_time":"2025-12-15T13:59:22.426797","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import gc\ndel dd\ngc.collect()","metadata":{"execution":{"iopub.execute_input":"2025-12-15T13:59:22.87542Z","iopub.status.busy":"2025-12-15T13:59:22.87517Z","iopub.status.idle":"2025-12-15T13:59:22.945259Z","shell.execute_reply":"2025-12-15T13:59:22.944672Z"},"papermill":{"duration":0.091924,"end_time":"2025-12-15T13:59:22.946373","exception":false,"start_time":"2025-12-15T13:59:22.854449","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Ensemble Everything\nBelow we ensemble 10x NN and 1x XGB. Note we also infer NN-v142 but we don't actually use it. To see the infer and train code for my strong NN, you need to read the code from the scripts in Kaggle inference script dataset [here][1]. The train code for all the NN is [here][2] in folder `train_scripts`. Each family of NN are the same NN but train with different time windows. The difference between families are different invariant input features. Read about all the NN in discussion [here][3]\n\nFamily one. Positions are agent centric coordinates:\n* NN-v258\n* NN-v211\n\nFamily two. Positions are distances relative to other body parts:\n* NN-v315\n* NN-v330\n* NN-v357\n* NN-v340\n* NN-v345\n\nFamily three. Positions are arena max min scaled:\n* NN-v417\n* NN-v421\n* NN-v422\n\n[1]: https://www.kaggle.com/datasets/cdeotte/mouse-comp-inference-scripts\n[2]: https://www.kaggle.com/datasets/cdeotte/mouse-comp-inference-scripts?select=train_scripts\n[3]: https://www.kaggle.com/competitions/MABe-mouse-behavior-detection/writeups/7th-place-gold-cnn-transformer-with-invariant-fe","metadata":{"papermill":{"duration":0.020612,"end_time":"2025-12-15T13:59:22.987934","exception":false,"start_time":"2025-12-15T13:59:22.967322","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# ====================== TEST-ONLY INFERENCE (LOW-RAM VERSION) ======================\nimport os, sys, json, math, gc, re, itertools, warnings, time, contextlib\nfrom typing import List, Dict, Tuple, Set\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport concurrent.futures as _fut\nimport threading as _th\n\n# ------------------------------- Config --------------------------------------\nTEST_PATH   = \"/kaggle/input/MABe-mouse-behavior-detection/\"\nTRAIN_PATH  = \"/kaggle/input/MABe-mouse-behavior-detection/\"\n\n# >>> NEW: two model lists for 2×GPU\nBUNDLE_PATH1 = [\"/kaggle/input/mouse-v142/mbnetdual_bundle_v142.pt\",\"/kaggle/input/mouse-v142/mbnetdual_bundle_v143.pt\",\"/kaggle/input/mouse-v142/mbnetdual_bundle_v144.pt\"]\nBUNDLE_PATH2 = [\"/kaggle/input/mouse-v142/mbnetdual_bundle_v149.pt\",\"/kaggle/input/mouse-v142/mbnetdual_bundle_v150.pt\",\"/kaggle/input/mouse-v142/mbnetdual_bundle_v151.pt\"]\n\n# Inference knobs\nWINDOW_T       = 64          # model context length\nSTEP_T         = WINDOW_T//2\nBATCH_VIDEOS   = 1           # videos per outer batch (smaller = lower RAM)\nBATCH_WINDOWS  = 1024*2      # temporal windows per forward pass (GPU VRAM control)\nSMOOTH_WIN     = 5           # rolling window for light smoothing\nDEFAULT_THR    = 0.7\nOUT_PATH       = \"submission.csv\"\nDEBUG          = 0           # 0=off; >0: use that many rows from train.csv (non-NA behaviors)\nENSEMBLE_N     = 10\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nTORCH_DTYPE = torch.float32\n\ndef log(msg): \n    print(msg, flush=True)\n\ndef _gpu_mem(dev: str | None = None):\n    if not torch.cuda.is_available(): \n        return \"gpu=NA\"\n    if dev is not None:\n        torch.cuda.synchronize(device=dev)\n    a = torch.cuda.memory_allocated() / (1024**3)\n    r = torch.cuda.memory_reserved() / (1024**3)\n    return f\"gpu={a:.1f}/{r:.1f} GB\"\n\nwarnings.filterwarnings(\"ignore\", category=pd.errors.PerformanceWarning)\nlog(f\"[boot] DEVICE={DEVICE}, torch={torch.__version__}, cuda={torch.cuda.is_available()}\")\n\n# --------------------------- Model definitions -------------------------------\nclass SinusoidalPE(nn.Module):\n    def __init__(self, d_model, max_len=200000):\n        super().__init__()\n        pe = torch.zeros(max_len, d_model)\n        pos = torch.arange(0, max_len).unsqueeze(1)\n        div = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))\n        pe[:, 0::2] = torch.sin(pos * div); pe[:, 1::2] = torch.cos(pos * div)\n        self.register_buffer(\"pe\", pe, persistent=False)\n    def forward(self, x):  # x: [B,T,D]\n        return x + self.pe[:x.size(1)]\n\nclass TCNBlock(nn.Module):\n    def __init__(self, c_in, c_out, k=5, d=1, p=0.2):\n        super().__init__()\n        pad = (k-1)//2 * d\n        self.conv1 = nn.Conv1d(c_in,  c_out, k, padding=pad, dilation=d)\n        self.conv2 = nn.Conv1d(c_out, c_out, k, padding=pad, dilation=d)\n        self.norm1 = nn.BatchNorm1d(c_out); self.norm2 = nn.BatchNorm1d(c_out)\n        self.drop  = nn.Dropout(p)\n        self.proj  = nn.Identity() if c_in==c_out else nn.Conv1d(c_in, c_out, 1)\n    def forward(self, x):  # [B,C,T]\n        r = self.proj(x)\n        x = self.drop(F.gelu(self.norm1(self.conv1(x))))\n        x = self.drop(F.gelu(self.norm2(self.conv2(x))))\n        return x + r\n\nclass DS_TCNEncoder(nn.Module):\n    def __init__(self, c_in, width=512, n_tcn=3, dropout=0.2):\n        super().__init__()\n        chans = [c_in, width // 2, width, width]\n        dil = [1, 2, 4][:max(1, n_tcn)]\n        blocks, c = [], c_in\n        for i, d in enumerate(dil):\n            c_out = chans[min(i + 1, len(chans) - 1)]\n            blocks.append(TCNBlock(c, c_out, k=5, d=d, p=dropout))\n            c = c_out\n        self.net = nn.Sequential(*blocks)\n        self.out_dim = c\n    def forward(self, x):  # x: [B,T,C]\n        x = x.transpose(1, 2)   # [B,C,T]\n        x = self.net(x)\n        return x.transpose(1, 2) # [B,T,C]\n\nclass CrossAttnBlock(nn.Module):\n    def __init__(self, d_model, nhead=8, dropout=0.2):\n        super().__init__()\n        self.ln_q = nn.LayerNorm(d_model)\n        self.ln_kv = nn.LayerNorm(d_model)\n        self.attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)\n        self.drop = nn.Dropout(dropout)\n        self.ln2 = nn.LayerNorm(d_model)\n        self.ff  = nn.Sequential(\n            nn.Linear(d_model, 4 * d_model), nn.GELU(), nn.Dropout(dropout),\n            nn.Linear(4 * d_model, d_model)\n        )\n    def forward(self, q, kv):  # [B,T,D]\n        qn, kvn = self.ln_q(q), self.ln_kv(kv)\n        h, _ = self.attn(qn, kvn, kvn, need_weights=False)\n        x = q + self.drop(h)\n        y = self.ff(self.ln2(x))\n        return x + self.drop(y)\n\nclass MouseBehaviorNetDualStream(nn.Module):\n    def __init__(self, in_dim: int, num_actions: int, feat_names: List[str],\n                 width: int = 512, n_tcn: int = 3, nhead: int = 8, dropout: float = 0.2):\n        super().__init__()\n        idx_agent, idx_target, idx_rel = [], [], []\n        for i, n in enumerate(map(str, feat_names)):\n            if n in (\"rel_dx\", \"rel_dy\", \"rel_dist\"):\n                idx_rel.append(i)\n            elif n.startswith(\"t_\") or n.startswith(\"m_t_\"):\n                idx_target.append(i)\n            else:\n                idx_agent.append(i)\n        if not idx_agent or not idx_target:\n            raise ValueError(\"Need unprefixed agent features and 't_*'/'m_t_*' target features.\")\n        self.register_buffer(\"idx_agent\", torch.tensor(idx_agent, dtype=torch.long), persistent=False)\n        self.register_buffer(\"idx_target\", torch.tensor(idx_target, dtype=torch.long), persistent=False)\n        self.register_buffer(\"idx_rel\",    torch.tensor(idx_rel,    dtype=torch.long), persistent=False)\n\n        d_agent, d_target, d_rel = len(idx_agent), len(idx_target), len(idx_rel)\n        shared_c = max(d_agent, d_target)\n        self.adapt_agent  = nn.Identity() if d_agent  == shared_c else nn.Linear(d_agent,  shared_c, bias=False)\n        self.adapt_target = nn.Identity() if d_target == shared_c else nn.Linear(d_target, shared_c, bias=False)\n\n        self.enc = DS_TCNEncoder(c_in=shared_c, width=width, n_tcn=n_tcn, dropout=dropout)\n        self.pe  = SinusoidalPE(self.enc.out_dim)\n        self.cross_at = CrossAttnBlock(self.enc.out_dim, nhead=nhead, dropout=dropout)  # agent attends target\n        self.cross_ta = CrossAttnBlock(self.enc.out_dim, nhead=nhead, dropout=dropout)  # target attends agent\n        self.rel_proj = nn.Linear(d_rel, self.enc.out_dim, bias=False) if d_rel > 0 else None\n\n        fuse_in = self.enc.out_dim * 2 + (self.enc.out_dim if self.rel_proj is not None else 0)\n        self.fuse = nn.Sequential(nn.LayerNorm(fuse_in), nn.Linear(fuse_in, width), nn.GELU(), nn.Dropout(dropout))\n        enc_layer = nn.TransformerEncoderLayer(\n            d_model=width, nhead=nhead, dim_feedforward=width * 4, dropout=dropout,\n            batch_first=True, activation=\"gelu\", norm_first=True\n        )\n        self.refine = nn.TransformerEncoder(enc_layer, num_layers=1)\n        self.head = nn.Sequential(nn.Linear(width, width), nn.GELU(), nn.Dropout(dropout), nn.Linear(width, num_actions))\n\n    def forward(self, x):  # [B,T,D]\n        Xa = x.index_select(-1, self.idx_agent)\n        Xt = x.index_select(-1, self.idx_target)\n        Xr = x.index_select(-1, self.idx_rel) if self.idx_rel.numel() else None\n        Xa = self.adapt_agent(Xa); Xt = self.adapt_target(Xt)\n        Ha = self.enc(Xa); Ht = self.enc(Xt)\n        Ha = self.pe(Ha); Ht = self.pe(Ht)\n        A  = self.cross_at(Ha, Ht); T  = self.cross_ta(Ht, Ha)\n        if Xr is not None and self.rel_proj is not None:\n            R = self.rel_proj(Xr); Z = torch.cat([A, T, R], dim=-1)\n        else:\n            Z = torch.cat([A, T], dim=-1)\n        Z = self.fuse(Z); Z = self.refine(Z)\n        return self.head(Z)\n\n# ----------------------------- Bundle I/O ------------------------------------\ndef _safe_torch_load(path: str):\n    try:\n        return torch.load(path, map_location=\"cpu\", weights_only=True)  # PyTorch ≥ 2.4\n    except TypeError:\n        return torch.load(path, map_location=\"cpu\")\n\ndef _infer_in_dim_from_state_dict(sd: dict) -> int:\n    preferred = [\"tcn.0.conv1.weight\", \"enc.net.0.conv1.weight\"]\n    for k in preferred:\n        if k in sd:\n            w = sd[k]; assert w.ndim in (3,4)\n            return w.shape[1]\n    cands = [k for k,v in sd.items() if k.endswith(\"conv1.weight\") and getattr(v, \"ndim\", 0) in (3,4)]\n    assert cands, \"Cannot infer in_dim\"\n    return sd[cands[0]].shape[1]\n\ndef _infer_num_actions_from_state_dict(sd: dict) -> int:\n    head_2d = [(k, v) for k, v in sd.items() if k.startswith(\"head.\") and getattr(v, \"ndim\", 0) == 2]\n    assert head_2d, \"Cannot infer num_actions\"\n    last_key, last_w = sorted(head_2d, key=lambda kv: kv[0])[-1]\n    return last_w.shape[0]\n\ndef load_inference_bundle(load_path: str, device: str = \"cuda\"):\n    bundle = _safe_torch_load(load_path)\n    cls_name = bundle.get(\"model_class\", \"MouseBehaviorNetDualStream\")\n    mk = dict(bundle.get(\"model_kwargs\", {}))\n    sd = bundle[\"model_state\"]\n    if mk.get(\"in_dim\") is None: mk[\"in_dim\"] = _infer_in_dim_from_state_dict(sd)\n    if mk.get(\"num_actions\") is None: mk[\"num_actions\"] = _infer_num_actions_from_state_dict(sd)\n    if cls_name == \"MouseBehaviorNetDualStream\" and \"feat_names\" not in mk:\n        mk[\"feat_names\"] = list(bundle.get(\"feat_cols\", []))\n    cls = globals().get(cls_name, None)\n    assert cls is not None, f\"Model class '{cls_name}' not defined\"\n    model = cls(**mk).to(device)\n    model.load_state_dict(sd, strict=True)\n    model.eval()\n\n    actions   = bundle[\"actions\"]\n    feat_cols = bundle[\"feat_cols\"]\n    scaler    = bundle[\"scaler\"]\n    #thr       = bundle.get(\"thresholds\", {})\n    thr       = {'allogroom': 0.34000000000000014, 'approach': 0.9400000000000001, 'attack': 0.77, 'attemptmount': 0.53, 'avoid': 0.75, 'biteobject': 0.55, 'chase': 0.8400000000000001, 'chaseattack': 0.76, 'climb': 0.48000000000000004, 'defend': 0.59, 'dig': 0.8, 'disengage': 0.8400000000000001, 'dominance': 0.63, 'dominancegroom': 0.66, 'dominancemount': 0.93, 'ejaculate': 0.35000000000000003, 'escape': 0.8300000000000001, 'exploreobject': 0.7400000000000001, 'flinch': 0.78, 'follow': 0.89, 'freeze': 0.30000000000000004, 'genitalgroom': 0.59, 'huddle': 0.6, 'intromit': 0.49000000000000005, 'mount': 0.7500000000000001, 'rear': 0.64, 'reciprocalsniff': 0.87, 'rest': 0.66, 'run': 0.89, 'selfgroom': 0.8, 'shepherd': 0.88, 'sniff': 0.61, 'sniffbody': 0.63, 'sniffface': 0.91, 'sniffgenital': 0.68, 'submit': 0.37000000000000005, 'tussle': 0.5499999999999999}\n    id_cols   = bundle.get(\"id_cols\", [\"video_id\",\"frame\",\"agent_id\",\"target_id\"])\n\n    log(f\"[bundle] class={cls_name} actions={len(actions)} feats={len(feat_cols)}\")\n    return model, actions, feat_cols, scaler, thr, id_cols\n\n# ----------------------------- Meta & I/O ------------------------------------\ndef _norm_cols(df: pd.DataFrame) -> pd.DataFrame:\n    df = df.copy()\n    df.columns = (df.columns.str.strip().str.lower()\n                  .str.replace(r\"[^\\w]+\",\"_\",regex=True)\n                  .str.replace(\"__+\",\"_\",regex=True)\n                  .str.strip(\"_\"))\n    return df\n\ndef read_csv_any(root: str, name: str) -> pd.DataFrame | None:\n    p = os.path.join(root, name)\n    if not os.path.exists(p): \n        return None\n    return _norm_cols(pd.read_csv(p))\n\ndef load_meta_row_any(lab_id: str, vid: int, prefer: str = \"train\") -> pd.Series:\n    lab_id, vid = str(lab_id).strip(), int(vid)\n    order = [(\"train\", TRAIN_PATH, \"train.csv\"), (\"test\", TEST_PATH, \"test.csv\")]\n    if prefer != \"train\":\n        order.reverse()\n    for _, root, csv in order:\n        meta = read_csv_any(root, csv)\n        if meta is None: \n            continue\n        hit = meta[(meta[\"lab_id\"].astype(str) == lab_id) & (pd.to_numeric(meta[\"video_id\"], errors=\"coerce\") == vid)]\n        if len(hit):\n            return hit.iloc[0]\n    raise KeyError(f\"No meta row for lab_id='{lab_id}', video_id={vid}\")\n\ndef tracking_path_any(lab: str, vid: int) -> str:\n    lab = str(lab).strip()\n    candidates = [\n        os.path.join(TEST_PATH,  \"test_tracking\",  lab, f\"{vid}.parquet\"),\n        os.path.join(TEST_PATH,  \"train_tracking\", lab, f\"{vid}.parquet\"),\n        os.path.join(TRAIN_PATH, \"test_tracking\",  lab, f\"{vid}.parquet\"),\n        os.path.join(TRAIN_PATH, \"train_tracking\", lab, f\"{vid}.parquet\"),\n    ]\n    for p in candidates:\n        if os.path.exists(p):\n            return p\n    raise FileNotFoundError(f\"Tracking parquet not found for lab='{lab}', vid={vid}\")\n\n# ----------------------------- Feature build ---------------------------------\nDROP_PARTS = [\n    'headpiece_bottombackleft','headpiece_bottombackright','headpiece_bottomfrontleft','headpiece_bottomfrontright',\n    'headpiece_topbackleft','headpiece_topbackright','headpiece_topfrontleft','headpiece_topfrontright',\n    'spine_1','spine_2','tail_middle_1','tail_middle_2','tail_midpoint'\n]\nALL_PARTS_FULL = [\n    \"ear_left\",\"ear_right\",\"hip_left\",\"hip_right\",\"neck\",\"nose\",\"tail_base\",\n    \"body_center\",\"lateral_left\",\"lateral_right\",\"tail_tip\",\"tail_midpoint\",\n    \"spine_1\",\"spine_2\",\"tail_middle_1\",\"tail_middle_2\",\"head\",\n    \"headpiece_bottombackleft\",\"headpiece_bottombackright\",\"headpiece_bottomfrontleft\",\"headpiece_bottomfrontright\",\n    \"headpiece_topbackleft\",\"headpiece_topbackright\",\"headpiece_topfrontleft\",\"headpiece_topfrontright\",\n]\nALL_PARTS = [p for p in ALL_PARTS_FULL if p not in DROP_PARTS]\n\ndef load_tracking_parquet(path: str) -> pd.DataFrame:\n    return pd.read_parquet(path)\n\ndef pivot_mouse_frame(tracking_df: pd.DataFrame, parts_order: List[str]) -> Dict[int, pd.DataFrame]:\n    mice = sorted(tracking_df[\"mouse_id\"].unique())\n    out = {}\n    # build only needed x/y columns to keep memory in check\n    for m in mice:\n        sub = tracking_df.loc[tracking_df[\"mouse_id\"]==m, [\"video_frame\",\"bodypart\",\"x\",\"y\"]]\n        wide = sub.pivot_table(index=\"video_frame\", columns=\"bodypart\", values=[\"x\",\"y\"])\n        cols = [(c0, p) for p in parts_order for c0 in (\"x\",\"y\")]\n        wide = wide.reindex(columns=pd.MultiIndex.from_tuples(cols, names=[\"coord\",\"bodypart\"]))\n        wide.columns = [f\"{c0}_{c1}\" for c0,c1 in wide.columns.to_list()]\n        out[m] = wide.sort_index()\n        del sub, wide\n    return out\n\ndef fill_and_normalize_with_masks(wide: pd.DataFrame, pix_per_cm: float, width_pix: float, height_pix: float) -> pd.DataFrame:\n    df = wide.copy()\n    ever_obs = ~df.isna().all(axis=0)\n    masks = pd.DataFrame({f\"m_{c}\": np.float32(ever_obs[c]) for c in df.columns}, index=df.index)\n    df = df.interpolate(limit_direction=\"both\").ffill().bfill().fillna(0.0)\n\n    # FIX: don't reference 'ppcm' before assignment; use the function arg 'pix_per_cm'\n    ppcm = float(pix_per_cm) if pd.notna(pix_per_cm) and float(pix_per_cm) != 0 else 1.0\n\n    cx, cy = float(width_pix)/2.0, float(height_pix)/2.0\n    x_cols = [c for c in df.columns if c.startswith(\"x_\")]\n    y_cols = [c for c in df.columns if c.startswith(\"y_\")]\n    if x_cols: df[x_cols] = (df[x_cols].astype(float) - cx) / ppcm\n    if y_cols: df[y_cols] = (df[y_cols].astype(float) - cy) / ppcm\n    df = df.replace([np.inf,-np.inf], np.nan).fillna(0.0)\n    return pd.concat([df, masks], axis=1)\n\n\ndef add_velocities(wide_with_masks: pd.DataFrame, fps: float) -> pd.DataFrame:\n    df = wide_with_masks.copy()\n    coord_cols = [c for c in df.columns if c.startswith(\"x_\") or c.startswith(\"y_\")]\n    vel = df[coord_cols].diff().fillna(0.0) * float(fps)\n    vel.columns = [f\"v{c}\" for c in coord_cols]\n    mask_names = [f\"m_{c}\" for c in coord_cols]\n    have_mask = [m for m in mask_names if m in df.columns]\n    mv_dict = {f\"m_v{m[2:]}\": df[m] for m in have_mask}\n    mv = pd.DataFrame(mv_dict, index=df.index) if mv_dict else pd.DataFrame(index=df.index)\n    out = pd.concat([df, vel, mv], axis=1)\n    del df, vel, mv\n    return out\n\n# ---- Robust parsing for behaviors_labeled (handles funny quoting/mixed forms)\n_TRIPLET_REGEX = re.compile(\n    r\"\"\"^\\s*\n        ['\"]?\\s*(?P<a>(?:mouse)?\\d+|self|same|agent)\\s*['\"]?\\s*,\\s*\n        ['\"]?\\s*(?P<t>(?:mouse)?\\d+|self|same|agent)\\s*['\"]?\\s*,\\s*\n        ['\"]?\\s*(?P<act>[A-Za-z_]+)\\s*['\"]?\n        \\s*$\"\"\", re.X | re.IGNORECASE\n)\ndef _safe_parse_behaviors_labeled(x):\n    if pd.isna(x): return set()\n    if isinstance(x, (list, tuple, set)): return {str(s).strip() for s in x if str(s).strip()}\n    s = str(x).strip()\n    if not s: return set()\n    try:\n        parsed = json.loads(s)\n        if isinstance(parsed, (list, tuple, set)):\n            return {str(t).strip() for t in parsed if str(t).strip()}\n        if isinstance(parsed, str):\n            st = parsed.strip()\n            return {st} if st else set()\n    except Exception:\n        pass\n    if \"|\" in s:\n        return {t.strip() for t in s.split(\"|\") if t.strip()}\n    return {s}\n\ndef _split_triplet_smart(trip: str):\n    if not isinstance(trip, str): return None\n    s = trip.strip().strip(\"[](){} \\t\\r\\n\").replace(\"’\", \"'\")\n    m = _TRIPLET_REGEX.match(s)\n    if m:\n        return m.group(\"a\").lower(), m.group(\"t\").lower(), m.group(\"act\").strip()\n    parts = [p.strip().strip(\"\\\"'[](){} \\t\\r\\n\") for p in s.split(\",\")]\n    if len(parts) < 3: return None\n    return parts[0].lower(), parts[1].lower(), parts[2].strip()\n\ndef _mouse_to_int(tok: str):\n    if not isinstance(tok, str): return None\n    s = tok.lower().replace(\" \", \"\").replace(\"mouse\", \"\")\n    m = re.search(r\"\\d+\", s)\n    return int(m.group(0)) if m else None\n\ndef _parse_triplet_to_numeric(triplet: str):\n    split = _split_triplet_smart(triplet)\n    if not split: return None\n    a_tok, t_tok, act = split\n    ai = None if a_tok in {\"self\",\"same\",\"agent\"} else _mouse_to_int(a_tok)\n    if ai is None: return None\n    if t_tok in {\"self\", \"same\", \"agent\"}:\n        ti = ai\n    else:\n        ti = _mouse_to_int(t_tok)\n        if ti is None: return None\n    act = str(act).strip()\n    if not act: return None\n    return (ai, ti, act)\n\ndef build_active_map_string(meta_df: pd.DataFrame) -> Dict[int, Set[str]]:\n    amap: Dict[int, Set[str]] = {}\n    if meta_df is None or \"video_id\" not in meta_df.columns or \"behaviors_labeled\" not in meta_df.columns:\n        return amap\n    for vid, g in meta_df.groupby(\"video_id\", sort=False):\n        S: Set[str] = set()\n        for raw in g[\"behaviors_labeled\"].dropna():\n            for trip in _safe_parse_behaviors_labeled(raw):\n                parsed = _parse_triplet_to_numeric(trip)\n                if parsed is None: \n                    continue\n                ai, ti, act = parsed\n                S.add(f\"{ai},{ti},{act}\")\n        amap[int(vid)] = S\n    return amap\n\n# -------------------------- Pairwise features (DS) ---------------------------\ndef _make_target_prefixed(df: pd.DataFrame) -> pd.DataFrame:\n    keep = [c for c in df.columns if c.startswith((\"x_\",\"y_\",\"vx_\",\"vy_\",\"m_x_\",\"m_y_\",\"m_vx_\",\"m_vy_\"))]\n    mapping = {c: (\"m_t_\" + c[2:] if c.startswith(\"m_\") else \"t_\" + c) for c in keep}\n    return df[keep].rename(columns=mapping)\n\ndef pairwise_features(per_mouse_wide: Dict[int, pd.DataFrame],\n                      allowed_pairs: List[Tuple[int,int]] | None = None) -> pd.DataFrame:\n    if not per_mouse_wide:\n        return pd.DataFrame(columns=[\"frame\",\"agent_id\",\"target_id\"])\n    mice = sorted(per_mouse_wide.keys())\n    mice_set = set(mice)\n\n    it = iter(mice)\n    first = next(it)\n    common_index = per_mouse_wide[first].index\n    for m in it:\n        common_index = common_index.intersection(per_mouse_wide[m].index)\n\n    per_mouse_aligned: Dict[int, pd.DataFrame] = {}\n    for m in mice:\n        dfm = per_mouse_wide[m].loc[common_index]\n        if len(dfm):\n            per_mouse_aligned[m] = dfm\n\n    if not per_mouse_aligned:\n        return pd.DataFrame(columns=[\"frame\",\"agent_id\",\"target_id\"])\n\n    if allowed_pairs is not None:\n        pairs = [(a, t) for (a, t) in allowed_pairs if a in per_mouse_aligned and t in per_mouse_aligned]\n        if not pairs:\n            avail = sorted(per_mouse_aligned.keys())\n            pairs = list(itertools.product(avail, avail))\n    else:\n        avail = sorted(per_mouse_aligned.keys())\n        pairs = list(itertools.product(avail, avail))\n\n    if not pairs:\n        return pd.DataFrame(columns=[\"frame\",\"agent_id\",\"target_id\"])\n\n    candidates = [p for p in [\"nose\",\"body_center\",\"neck\",\"tail_base\",\"ear_left\",\"ear_right\"] if p in ALL_PARTS]\n\n    rows = []\n    for agent, target in pairs:\n        A = per_mouse_aligned.get(agent)\n        T = per_mouse_aligned.get(target)\n        if A is None or T is None or A.empty or T.empty:\n            continue\n\n        feat = pd.concat([A, _make_target_prefixed(T)], axis=1)\n        picked = None\n        for p in candidates:\n            Ax, Ay = f\"x_{p}\", f\"y_{p}\"\n            if Ax in A and Ay in A and Ax in T and Ay in T:\n                picked = p\n                break\n\n        if picked is not None:\n            ax = A[f\"x_{picked}\"].to_numpy(); ay = A[f\"y_{picked}\"].to_numpy()\n            tx = T[f\"x_{picked}\"].to_numpy(); ty = T[f\"y_{picked}\"].to_numpy()\n            dx = tx - ax; dy = ty - ay\n            feat = feat.assign(rel_dx=dx, rel_dy=dy, rel_dist=np.sqrt(dx*dx + dy*dy))\n        else:\n            feat = feat.assign(rel_dx=0.0, rel_dy=0.0, rel_dist=0.0)\n\n        feat[\"agent_id\"] = agent\n        feat[\"target_id\"] = target\n        rows.append(feat)\n\n    if not rows:\n        return pd.DataFrame(columns=[\"frame\",\"agent_id\",\"target_id\"])\n\n    out = pd.concat(rows, copy=False).reset_index(names=\"frame\")\n    id_cols = [\"frame\",\"agent_id\",\"target_id\"]\n    feat_cols = [c for c in out.columns if c not in id_cols]\n    return out[id_cols + feat_cols]\n\ndef build_video_feats_test_only(lab_id: str, vid: int,\n                                allowed_pairs: List[Tuple[int,int]] | None = None) -> pd.DataFrame:\n    meta = load_meta_row_any(lab_id, vid, prefer=\"train\")\n    fps = float(meta.get(\"frames_per_second\", 30.0))\n    ppcm = meta.get(\"pix_per_cm_approx\", np.nan)\n    ppcm = float(ppcm) if pd.notna(ppcm) and float(ppcm) != 0 else np.nan\n    width_pix  = float(meta.get(\"video_width_pix\", np.nan))\n    height_pix = float(meta.get(\"video_height_pix\", np.nan))\n    if (not pd.notna(ppcm)) and all(pd.notna(meta.get(c, np.nan)) for c in [\"arena_width_cm\",\"arena_height_cm\"]):\n        aw = float(meta[\"arena_width_cm\"]); ah = float(meta[\"arena_height_cm\"])\n        if aw>0 and ah>0 and pd.notna(width_pix) and pd.notna(height_pix):\n            ppcm = ((width_pix/aw) + (height_pix/ah)) / 2.0\n    if not pd.notna(ppcm) or ppcm == 0: ppcm = 1.0\n\n    path = tracking_path_any(lab_id, vid)\n    track = load_tracking_parquet(path)\n    track = track[~track[\"bodypart\"].isin(DROP_PARTS)].copy()\n\n    per_mouse = pivot_mouse_frame(track, ALL_PARTS)\n    del track\n    for m in list(per_mouse.keys()):\n        pm = fill_and_normalize_with_masks(per_mouse[m], ppcm, width_pix, height_pix)\n        per_mouse[m] = add_velocities(pm, fps)\n        del pm\n\n    if allowed_pairs is not None:\n        mice_present = set(per_mouse.keys())\n        allowed_pairs = [(a, t) for (a, t) in allowed_pairs if a in mice_present and t in mice_present]\n        if not allowed_pairs:\n            allowed_pairs = None  # fall back to all x all\n\n    feats = pairwise_features(per_mouse, allowed_pairs=allowed_pairs)\n    del per_mouse\n    feats[\"video_id\"] = int(vid)\n    return feats\n\n# ------------------------------- Masking -------------------------------------\ndef mask_probs_numpy_rle(probs: pd.DataFrame, ACTIONS: List[str], active_map: Dict[int, Set[str]],\n                         copy=True) -> pd.DataFrame:\n    df = probs.copy() if copy else probs\n    if not len(df): return df\n    action_cols = [f\"action_{a}\" for a in ACTIONS]\n    act_block = df[action_cols].to_numpy(copy=False)\n    N, A = act_block.shape\n    vid = df[\"video_id\"].to_numpy(np.int64, copy=False)\n    ag  = df[\"agent_id\"].to_numpy(np.int64, copy=False)\n    tg  = df[\"target_id\"].to_numpy(np.int64, copy=False)\n\n    act_pos = {a: i for i, a in enumerate(ACTIONS)}\n    allow: Dict[int, Dict[Tuple[int,int], np.ndarray]] = {}\n    for v, triples in active_map.items():\n        v = int(v)\n        d = allow.setdefault(v, {})\n        for s in triples:\n            sag, stg, sa = s.split(\",\")\n            key = (int(sag), int(stg))\n            arr = d.get(key)\n            if arr is None:\n                arr = np.zeros(A, dtype=bool)\n                d[key] = arr\n            i = act_pos.get(sa)\n            if i is not None:\n                arr[i] = True\n\n    if N == 1:\n        starts = np.array([0], dtype=np.int64); ends = np.array([1], dtype=np.int64)\n    else:\n        change = (vid[1:] != vid[:-1]) | (ag[1:] != ag[:-1]) | (tg[1:] != tg[:-1])\n        boundaries = np.flatnonzero(change) + 1\n        starts = np.concatenate(([0], boundaries))\n        ends   = np.concatenate((boundaries, [N]))\n\n    for s, e in zip(starts, ends):\n        v, a_, t_ = int(vid[s]), int(ag[s]), int(tg[s])\n        d = allow.get(v)\n        if d is None:\n            act_block[s:e, :] = 0.0; continue\n        mask = d.get((a_, t_))\n        if mask is None:\n            act_block[s:e, :] = 0.0; continue\n        if not mask.all():\n            disallowed = ~mask\n            act_block[s:e, disallowed] = 0.0\n\n    df[action_cols] = act_block\n    return df\n\n# ----------------------------- Intervals -------------------------------------\nfrom typing import List, Dict\nimport numpy as np\nimport pandas as pd\n\ndef probs_to_nonoverlapping_intervals(\n    prob_df: pd.DataFrame,\n    actions: List[str],\n    min_len: int = 3,\n    max_gap: int = 2,\n    lab: str | None = None,\n    tie_config: Dict[str, dict] | None = None,\n) -> pd.DataFrame:\n    \"\"\"\n    Convert frame-level probs to non-overlapping intervals for a *single lab*.\n\n    - Thresholds are taken from TT_PER_LAB_NN[lab] and TT_PER_LAB_XGB[lab].\n    - If `tie_config` is given and contains `lab`, we apply per-lab tie manipulation\n      on frames where multiple actions pass threshold.\n\n    tie_config[lab] format (per what you've been using in validation):\n      {\n        \"boost\":    { \"action_name\": delta, ... },\n        \"penalize\": { \"action_name\": delta, ... },\n        \"prefer\":   [\n            (\"winner_action\", \"loser_action\", margin),\n            ...\n        ],\n      }\n    \"\"\"\n    out: list[dict] = []\n    act_cols = [f\"action_{a}\" for a in actions]\n\n    # Per-lab thresholds (you’re overriding per_action_thresh here anyway)\n    per_action_thresh = TT_PER_LAB_NN[lab]\n    print(f\"Using NN lab = {lab} with th = {per_action_thresh}\")\n    per_action_thresh2 = TT_PER_LAB_XGB[lab]\n    print(f\"Using XGB lab = {lab} with th = {per_action_thresh2}\")\n\n    thr2 = np.array(\n        [per_action_thresh2.get(a, 0.2) if per_action_thresh2 else 0.2 for a in actions],\n        dtype=np.float32,\n    )\n    thr = np.array(\n        [per_action_thresh.get(a, 0.75) if per_action_thresh else 0.75 for a in actions],\n        dtype=np.float32,\n    )\n\n    # Optional per-lab tie rules\n    lab_tie_cfg = None\n    if tie_config is not None and lab is not None and lab in tie_config:\n        lab_tie_cfg = tie_config[lab]\n        # map action name -> column index\n        action_to_idx = {a: i for i, a in enumerate(actions)}\n    else:\n        action_to_idx = {}\n\n    for (vid, ag, tg), grp in prob_df.groupby([\"video_id\", \"agent_id\", \"target_id\"], sort=False):\n        g = grp.sort_values(\"frame\")\n        frames = g[\"frame\"].to_numpy()\n        P = g[act_cols].to_numpy(np.float32)  # [T, num_actions]\n\n        # blend of NN + XGB thresholds\n        blended_thr = (1.0 - XGB_WGT) * thr[None, :] + XGB_WGT * thr2[None, :]\n        pass_mask = (P >= blended_thr)\n\n        # ---- Tie manipulation (per lab) ----\n        P_adj = P.copy()\n        if lab_tie_cfg is not None:\n            # frames where 2+ actions pass threshold\n            multi_mask = (pass_mask.sum(axis=1) > 1)\n\n            if multi_mask.any():\n                # 1) boosts\n                boost_cfg = lab_tie_cfg.get(\"boost\", {})\n                for act, delta in boost_cfg.items():\n                    idx = action_to_idx.get(act, None)\n                    if idx is not None:\n                        P_adj[multi_mask, idx] += float(delta)\n\n                # 2) penalties\n                penalize_cfg = lab_tie_cfg.get(\"penalize\", {})\n                for act, delta in penalize_cfg.items():\n                    idx = action_to_idx.get(act, None)\n                    if idx is not None:\n                        P_adj[multi_mask, idx] -= float(delta)\n\n                # 3) explicit preferences: (\"winner\", \"loser\", margin)\n                prefer_cfg = lab_tie_cfg.get(\"prefer\", [])\n                for winner_act, loser_act, margin in prefer_cfg:\n                    wi = action_to_idx.get(winner_act, None)\n                    li = action_to_idx.get(loser_act, None)\n                    if wi is None or li is None:\n                        continue\n                    # Only when both pass threshold on that frame\n                    fm = multi_mask & pass_mask[:, wi] & pass_mask[:, li]\n                    if fm.any():\n                        # Nudge winner upwards on these ambiguous frames\n                        P_adj[fm, wi] += float(margin)\n\n                # keep probabilities in a sane range\n                np.clip(P_adj, 0.0, 1.0, out=P_adj)\n\n        # --- Best-label decoding with thresholds + adjusted scores ---\n        P_masked = np.where(pass_mask, P_adj, -np.inf)\n        best_idx = np.argmax(P_masked, axis=1)\n        best_val = P_masked[np.arange(len(P_masked)), best_idx]\n        label = np.where(np.isfinite(best_val), best_idx, -1)\n\n        # --- Fill short gaps up to max_gap (unchanged) ---\n        if max_gap > 0:\n            i = 0\n            while i < len(label):\n                if label[i] >= 0:\n                    j = i\n                    while j + 1 < len(label) and label[j + 1] == label[i]:\n                        j += 1\n                    k = j + 1\n                    while k < len(label) and label[k] == -1:\n                        k += 1\n                    if k < len(label) and label[k] == label[i] and (k - j - 1) <= max_gap:\n                        label[j + 1:k] = label[i]\n                        j = k\n                    i = j + 1\n                else:\n                    i += 1\n\n        # --- Convert label sequence to intervals (unchanged) ---\n        def flush(s, e, idx):\n            if s is None:\n                return\n            if frames[e] - frames[s] + 1 >= min_len:\n                out.append(\n                    {\n                        \"video_id\": int(vid),\n                        \"agent_id\": int(ag),\n                        \"target_id\": int(tg),\n                        \"action\": actions[idx],\n                        \"start_frame\": int(frames[s]),\n                        \"stop_frame\": int(frames[e] + 1),\n                    }\n                )\n\n        s = None\n        cur = -1\n        for i_t, idx in enumerate(label):\n            if idx != cur:\n                if cur >= 0:\n                    flush(s, i_t - 1, cur)\n                s = i_t if idx >= 0 else None\n                cur = idx\n        if cur >= 0:\n            flush(s, len(label) - 1, cur)\n\n    return pd.DataFrame(out)\n\n# ------------------------------- Utils ---------------------------------------\ndef ensure_feat_alignment(df: pd.DataFrame, id_cols: List[str], FEAT_COLS: List[str]) -> pd.DataFrame:\n    for c in FEAT_COLS:\n        if c not in df.columns:\n            df[c] = 0.0\n    out = df[id_cols + FEAT_COLS].copy()\n    for c in (\"video_id\",\"frame\",\"agent_id\",\"target_id\"):\n        out[c] = pd.to_numeric(out[c], errors=\"coerce\").fillna(0).astype(np.int64)\n    for c in FEAT_COLS:\n        out[c] = pd.to_numeric(out[c], errors=\"coerce\").fillna(0).astype(np.float32)\n    return out\n\ndef robust_scale_inplace(df: pd.DataFrame, FEAT_COLS: List[str], scaler: dict):\n    med = pd.Series(scaler.get(\"med\", {}), dtype=np.float32)\n    iqr = pd.Series(scaler.get(\"iqr\", {}), dtype=np.float32).replace(0, 1.0)\n    low = float(scaler.get(\"clip_low\", -5.0)); high = float(scaler.get(\"clip_high\", 5.0))\n    X = df[FEAT_COLS].astype(np.float32)\n    X = (X - med.reindex(FEAT_COLS).fillna(0.0).astype(np.float32)) / iqr.reindex(FEAT_COLS).fillna(1.0).astype(np.float32)\n    X = X.clip(low, high).replace([np.inf,-np.inf], np.nan).fillna(0.0).astype(np.float32)\n    df.loc[:, FEAT_COLS] = X.values\n\n@torch.inference_mode()\ndef predict_groups(model: nn.Module, actions: List[str], df_part: pd.DataFrame,\n                   T=WINDOW_T, step=STEP_T, batch_windows=BATCH_WINDOWS, device: str = None) -> pd.DataFrame:\n    \"\"\"Stream groups to reduce RAM; stack up to `batch_windows` temporal windows per pass.\n       NEW: optional `device` to run on specific GPU in threaded mode.\"\"\"\n    use_device = device if device is not None else DEVICE\n    id_cols = [\"video_id\",\"frame\",\"agent_id\",\"target_id\"]\n    feat_cols = [c for c in df_part.columns if c not in id_cols]\n    act_cols = [f\"action_{a}\" for a in actions]\n\n    gobj = df_part.groupby([\"video_id\",\"agent_id\",\"target_id\"], sort=False)\n    n_groups = gobj.ngroups\n    log(f\"[predict] groups={n_groups} | T={T} step={step} stack={batch_windows} | { _gpu_mem(use_device) }\")\n\n    out_parts = []\n    gi = 0\n    for (vid,a,t), grp in gobj:\n        gi += 1\n        g = grp.sort_values(\"frame\")\n        X = g[feat_cols].to_numpy(np.float32, copy=False)\n        F = len(g)\n\n        preds  = np.zeros((F, len(actions)), np.float32)\n        counts = np.zeros((F, 1), np.float32)\n\n        starts = list(range(0, F, step))\n        wi = 0\n        while wi < len(starts):\n            this = starts[wi:wi+batch_windows]\n            if not this: break\n\n            lens = []\n            maxT = 0\n            for s in this:\n                e = min(s + T, F)\n                lens.append((s,e))\n                maxT = max(maxT, e - s)\n\n            batch = np.empty((len(lens), maxT, X.shape[1]), np.float32)\n            for i,(s,e) in enumerate(lens):\n                L = e - s\n                batch[i, :L] = X[s:e]\n                if L < maxT:\n                    batch[i, L:maxT] = X[e-1:e]\n\n            tb = torch.from_numpy(batch).to(use_device, non_blocking=True)\n            ctx = torch.autocast(device_type=\"cuda\", dtype=torch.float16) if torch.cuda.is_available() else contextlib.nullcontext()\n            with ctx:\n                out = model(tb).sigmoid().detach().cpu().numpy().astype(np.float32, copy=False)\n\n            for i,(s,e) in enumerate(lens):\n                L = e - s\n                preds[s:e]  += out[i, :L]\n                counts[s:e] += 1.0\n\n            del batch, tb, out, lens\n            wi += batch_windows\n\n        preds /= np.maximum(counts, 1.0)\n        part = pd.DataFrame(preds, columns=act_cols)\n        part[\"video_id\"]=int(vid); part[\"agent_id\"]=int(a); part[\"target_id\"]=int(t)\n        part[\"frame\"] = g[\"frame\"].to_numpy(copy=False)\n        out_parts.append(part)\n\n        if gi == 1 or gi == n_groups or (gi % 25 == 0):\n            log(f\"[predict] {gi}/{n_groups} (vid={vid}, a={a}, t={t}) | { _gpu_mem(use_device) }\")\n        del g, X, preds, counts, part, grp\n        if gi % 200 == 0:\n            gc.collect()\n\n    res = pd.concat(out_parts, ignore_index=True, copy=False) if out_parts else df_part.head(0)\n    del out_parts\n    gc.collect()\n    return res\n\n# ---------------------- NEW: multi-model helpers (2×GPU) ---------------------\nclass _DeviceEnsemble:\n    def __init__(self, device: str, bundle_paths: List[str]):\n        self.device = device\n        self.paths = [p for p in (bundle_paths or []) if isinstance(p, str)]\n        self.models: List[nn.Module] = []\n        self.scalers: List[dict] = []\n        self.thresholds_merged: Dict[str, float] = {}\n        self.first_thresholds: Dict[str, float] | None = None   # <<< NEW\n        self.ACTIONS: List[str] = []\n        self.FEAT_COLS: List[str] = []\n        self.id_cols: List[str] = []\n        self.lock = _th.Lock()\n\n    def load(self):\n        if torch.cuda.is_available() and self.device.startswith(\"cuda\"):\n            idx = int(self.device.split(\":\")[1])\n            torch.cuda.set_device(idx)\n        for i, p in enumerate(self.paths):\n            m, A, F, sc, thr, ids = load_inference_bundle(p, device=self.device)\n            self.models.append(m)\n            self.scalers.append(sc)\n            if i == 0:\n                self.ACTIONS = A; self.FEAT_COLS = F; self.id_cols = ids\n                if thr:                        # <<< NEW: remember thresholds from the first model on this device\n                    self.first_thresholds = dict(thr)\n            else:\n                if A != self.ACTIONS:\n                    log(f\"[warn] ACTIONS mismatch on {p}; using first bundle's order\")\n                if F != self.FEAT_COLS:\n                    log(f\"[warn] FEAT_COLS mismatch on {p}; columns will be aligned by name\")\n                if ids != self.id_cols:\n                    log(f\"[warn] id_cols mismatch on {p}; using first bundle's\")\n            if thr:\n                self.thresholds_merged.update(thr)\n\n    def predict_avg(self, feats: pd.DataFrame, T=WINDOW_T, step=STEP_T, batch_windows=BATCH_WINDOWS) -> pd.DataFrame:\n        \"\"\"Average logits/probs across this device's models.\"\"\"\n        if not self.models:\n            return pd.DataFrame(columns=[\"video_id\",\"agent_id\",\"target_id\",\"frame\"] + [f\"action_{a}\" for a in self.ACTIONS])\n        parts = []\n        for mi, model in enumerate(self.models, 1):\n            out = predict_groups(model, self.ACTIONS, feats, T=T, step=step, batch_windows=batch_windows, device=self.device)\n            parts.append(out)\n            gc.collect()\n        if len(parts) == 1:\n            return parts[0]\n        # align on keys then mean action columns\n        key = [\"video_id\",\"agent_id\",\"target_id\",\"frame\"]\n        merged = parts[0]\n        for k in range(1, len(parts)):\n            merged = merged.merge(parts[k], on=key, how=\"inner\", suffixes=(None, f\"__m{k}\"))\n        act_cols = [f\"action_{a}\" for a in self.ACTIONS]\n        # collect same-named cols + any suffixed copies\n        all_act_cols = [c for c in merged.columns if c.startswith(\"action_\")]\n        # average per action across model outputs\n        for a in act_cols:\n            like = [c for c in all_act_cols if c.split(\"__\")[0] == a]\n            merged[a] = merged[like].mean(axis=1).astype(np.float32)\n        # keep only one copy\n        keep = key + act_cols\n        return merged[keep]\n\ndef _weighted_average_two_probs(probs1: pd.DataFrame | None, n1: int,\n                                probs2: pd.DataFrame | None, n2: int,\n                                ACTIONS: List[str]) -> pd.DataFrame:\n    \"\"\"Average two prob tables with possibly different presence; weights = counts of models.\"\"\"\n    key = [\"video_id\",\"agent_id\",\"target_id\",\"frame\"]\n    act_cols = [f\"action_{a}\" for a in ACTIONS]\n    if (probs1 is None or probs1.empty) and (probs2 is None or probs2.empty):\n        return pd.DataFrame(columns=key + act_cols)\n    if probs1 is None or probs1.empty:\n        return probs2.copy()\n    if probs2 is None or probs2.empty:\n        return probs1.copy()\n    m = probs1.merge(probs2, on=key, how=\"inner\", suffixes=(\"_a\",\"_b\"))\n    for a in act_cols:\n        m[a] = ((m[f\"{a}_a\"] * float(max(n1,0))) + (m[f\"{a}_b\"] * float(max(n2,0)))) / float(max(n1,0) + max(n2,0))\n    return m[key + act_cols].astype({a: np.float32 for a in act_cols})\n\ndef smooth_probs_inplace(probs: pd.DataFrame, actions: List[str], win: int = SMOOTH_WIN):\n    if win is None or win <= 1 or probs.empty: \n        return\n    act_cols = [f\"action_{a}\" for a in actions]\n    probs.sort_values([\"video_id\",\"agent_id\",\"target_id\",\"frame\"], inplace=True)\n    pad = win // 2\n    kernel = np.ones(win, dtype=np.float32) / win\n    gobj = probs.groupby([\"video_id\",\"agent_id\",\"target_id\"], sort=False)\n    for _, idx in gobj.groups.items():\n        idx = np.asarray(idx)\n        arr = probs.loc[idx, act_cols].to_numpy(np.float32, copy=False)\n        for j in range(arr.shape[1]):\n            v = arr[:, j]\n            if v.size:\n                vp = np.pad(v, (pad, pad), mode=\"edge\")\n                arr[:, j] = np.convolve(vp, kernel, mode=\"valid\").astype(np.float32, copy=False)\n        probs.loc[idx, act_cols] = arr\n\n# --------------------------------- Main --------------------------------------\ndef main():\n    t0_all = time.time()\n\n    # Load meta/test list (same as before)\n    if DEBUG and DEBUG > 0:\n        train_meta = read_csv_any(TRAIN_PATH, \"train.csv\")\n        if train_meta is None:\n            raise FileNotFoundError(\"train.csv not found for DEBUG mode\")\n        train_meta = train_meta.loc[~train_meta[\"behaviors_labeled\"].isna()].reset_index()\n        rows = train_meta.index.to_list()[:DEBUG]\n        sel = train_meta.iloc[:DEBUG].copy()\n        TEST_LIST = list(sel.drop_duplicates(subset=[\"lab_id\",\"video_id\"], keep=\"first\")[[\"lab_id\",\"video_id\"]]\n                         .itertuples(index=False, name=None))\n        active_map_str = build_active_map_string(sel)\n        log(f\"[DEBUG] Using {len(sel)} rows from train.csv (behaviors_labeled notna). Row indices: {rows}\")\n        log(f\"[DEBUG] Unique (lab, video) count: {len(TEST_LIST)}\")\n        prefer_meta = \"train\"\n        del train_meta, sel\n    else:\n        test_meta = read_csv_any(TEST_PATH, \"test.csv\")\n        if test_meta is None:\n            raise FileNotFoundError(\"test.csv not found\")\n        TEST_LIST = list(test_meta.drop_duplicates(subset=[\"lab_id\",\"video_id\"], keep=\"first\")[[\"lab_id\",\"video_id\"]]\n                         .itertuples(index=False, name=None))\n        active_map_str = build_active_map_string(test_meta)\n        prefer_meta = \"train\"\n        del test_meta\n\n    # Preflight for tracking/meta availability\n    GOOD, BAD = [], []\n    for lab, vid in TEST_LIST:\n        try:\n            _ = tracking_path_any(lab, vid)\n            _ = load_meta_row_any(lab, vid, prefer=prefer_meta)\n            GOOD.append((lab, vid))\n        except Exception as e:\n            BAD.append((lab, vid, str(e)))\n    if BAD:\n        log(f\"[preflight] skipping {len(BAD)} videos due to missing meta/parquet\")\n    TEST_LIST = GOOD if GOOD else TEST_LIST\n    del GOOD, BAD\n    gc.collect()\n\n    # Prepare output\n    pd.DataFrame(columns=[\"row_id\",\"video_id\",\"agent_id\",\"target_id\",\"action\",\"start_frame\",\"stop_frame\"]).to_csv(OUT_PATH, index=False)\n\n    # >>> NEW: build per-device ensembles once (models pinned to each GPU)\n    ngpu = torch.cuda.device_count()\n    dev0 = \"cuda:0\" if torch.cuda.is_available() and ngpu >= 1 else \"cpu\"\n    dev1 = \"cuda:1\" if torch.cuda.is_available() and ngpu >= 2 else None\n\n    ens0 = _DeviceEnsemble(dev0, BUNDLE_PATH1); ens0.load()\n    ACTIONS = ens0.ACTIONS; FEAT_COLS = ens0.FEAT_COLS; id_cols = ens0.id_cols\n\n    ens1 = None\n    if dev1 is not None and BUNDLE_PATH2:\n        ens1 = _DeviceEnsemble(dev1, BUNDLE_PATH2); ens1.load()\n        if not ACTIONS:  ACTIONS  = ens1.ACTIONS\n        if not FEAT_COLS: FEAT_COLS = ens1.FEAT_COLS\n        if not id_cols:   id_cols   = ens1.id_cols\n\n    if not ACTIONS or not FEAT_COLS or not id_cols:\n        raise RuntimeError(\"Could not infer ACTIONS/FEAT_COLS/id_cols from bundles\")\n\n    # >>> NEW: thresholds from the FIRST available model only (ens0's first, else ens1's first)\n    if ens0.first_thresholds is not None:\n        thresholds = dict(ens0.first_thresholds)\n        thresholds_source = \"BUNDLE_PATH1[0]\"\n    elif ens1 is not None and ens1.first_thresholds is not None:\n        thresholds = dict(ens1.first_thresholds)\n        thresholds_source = \"BUNDLE_PATH2[0]\"\n    else:\n        thresholds = {}\n        thresholds_source = \"default (none in bundles)\"\n    log(f\"[bundle] thresholds source: {thresholds_source}\")\n\n    # Process in multi-video batches\n    log(f\"[infer] videos 1-{len(TEST_LIST)} / {len(TEST_LIST)} (batch={BATCH_VIDEOS}, win_batch={BATCH_WINDOWS}) | { _gpu_mem() }\")\n    row_counter = 0\n    total_pred_rows = 0\n\n    for bi in range(0, len(TEST_LIST), BATCH_VIDEOS):\n        t_batch = time.time()\n        batch = TEST_LIST[bi:bi+BATCH_VIDEOS]\n        tag = f\"[batch {bi//BATCH_VIDEOS + 1}/{(len(TEST_LIST)+BATCH_VIDEOS-1)//BATCH_VIDEOS}]\"\n        labs_vids_str = \" \".join([f\"({lab},{vid})\" for lab,vid in batch])\n        log(f\"{tag} building feats for {len(batch)} vids: {labs_vids_str}\")\n\n        # ---------- Single-thread: build feats for this batch ----------\n        feats_list = []\n        for lab, vid in batch:\n            allowed_set = active_map_str.get(int(vid), set())\n            pairs = sorted({(int(s.split(\",\")[0]), int(s.split(\",\")[1])) for s in allowed_set}) if allowed_set else None\n            f = build_video_feats_test_only(lab, vid, allowed_pairs=pairs)\n            feats_list.append(f)\n            log(f\"{tag} built feats vid={vid} rows={len(f)}\")\n        feats = pd.concat(feats_list, ignore_index=True, copy=False)\n        del feats_list\n        gc.collect()\n\n        total_pred_rows += len(feats)\n        log(f\"{tag} concat feats rows={len(feats)} | { _gpu_mem() }\")\n\n        # Align, scale (use ens0's scaler if available, else ens1's)\n        feats = ensure_feat_alignment(feats, id_cols, FEAT_COLS)\n        base_scaler = ens0.scalers[0] if ens0.scalers else (ens1.scalers[0] if ens1 and ens1.scalers else None)\n        if base_scaler is None:\n            raise RuntimeError(\"No scaler found in any bundle\")\n        robust_scale_inplace(feats, FEAT_COLS, base_scaler)\n\n        # ---------- Two threads: run per-device ensemble inference ----------\n        def _run_on_ens(ens: _DeviceEnsemble):\n            if ens is None or not ens.models:\n                return None\n            return ens.predict_avg(feats, T=WINDOW_T, step=STEP_T, batch_windows=BATCH_WINDOWS)\n\n        with _fut.ThreadPoolExecutor(max_workers=2) as ex:\n            futs = []\n            futs.append(ex.submit(_run_on_ens, ens0))\n            if ens1 is not None:\n                futs.append(ex.submit(_run_on_ens, ens1))\n            results = [f.result() for f in futs]\n\n        probs0 = results[0] if len(results) >= 1 else None\n        probs1 = results[1] if len(results) >= 2 else None\n        n0 = len(ens0.models)\n        n1 = len(ens1.models) if (ens1 is not None) else 0\n\n        # Weighted combine across devices\n        probs = _weighted_average_two_probs(probs0, n0, probs1, n1, ACTIONS)\n        log(f\"{tag} predict done | probs_rows={len(probs)} | { _gpu_mem() }\")\n        del probs0, probs1\n        gc.collect()\n\n       #############################\n        # -------- Ensemble with per-video probs from model1..modelN --------\n        key_cols = [\"video_id\", \"agent_id\", \"target_id\", \"frame\"]\n\n        # Ensure key dtypes are consistent (usually ints) in base probs\n        for c in key_cols:\n            probs[c] = pd.to_numeric(probs[c], errors=\"coerce\").astype(\"Int64\")\n\n        # Collect the corresponding probs for all videos in this batch\n        vids_in_batch = probs[\"video_id\"].dropna().unique().tolist()\n\n        # Start with base model (current probs) as the running ensemble\n        probs_ens = probs.copy()\n        CC = [c for c in probs_ens.columns if \"action\" in c]\n        probs_ens[CC] = 0.0\n        n_models = 0  # current ensemble has 0 models; base probs are zeroed out\n\n        # Loop over external models model1, model2, ..., model{ENSEMBLE_N}\n        for k in range(1, ENSEMBLE_N + 1):\n            model_dir = f\"/tmp/model{k}\"\n\n            # Load all per-video parquet files for this batch for model k\n            model_dfs = []\n            for vid in vids_in_batch:\n                vid_int = int(vid)\n                path = os.path.join(model_dir, f\"vid{vid_int}_probs.parquet\")\n                if not os.path.exists(path):\n                    raise FileNotFoundError(\n                        f\"Ensemble model folder '{model_dir}' probs file not found \"\n                        f\"for video_id={vid_int}: {path}\"\n                    )\n                df_m = pd.read_parquet(path)\n                # Make sure keys are same dtype\n                for c in key_cols:\n                    df_m[c] = pd.to_numeric(df_m[c], errors=\"coerce\").astype(\"Int64\")\n                model_dfs.append(df_m)\n\n            if not model_dfs:\n                raise RuntimeError(\n                    f\"No probs loaded from '{model_dir}' for batch videos: {vids_in_batch}\"\n                )\n\n            probs_m = pd.concat(model_dfs, ignore_index=True)\n            del model_dfs\n            gc.collect()\n\n            # Merge current ensemble with this model's probs on exact key\n            # Keep existing columns unchanged, suffix the new model's overlapping columns\n            probs_merged = probs_ens.merge(\n                probs_m,\n                on=key_cols,\n                how=\"inner\",\n                suffixes=(\"\", f\"_m{k}\"),\n            )\n            del probs_m\n            gc.collect()\n\n            if len(probs_merged) != len(probs_ens):\n                log(\n                    f\"{tag} WARNING: after merging model{k}, rows {len(probs_merged)} \"\n                    f\"!= previous {len(probs_ens)} \"\n                    f\"(some rows may be missing in model{k} files)\"\n                )\n\n            # Update running equal-weight average for each action column\n            # new_mean = (old_mean * n_models + new_model) / (n_models + 1)\n            for a in ACTIONS:\n                base = f\"action_{a}\"\n                col_new = f\"{base}_m{k}\"\n\n                if base not in probs_merged.columns or col_new not in probs_merged.columns:\n                    raise KeyError(\n                        f\"Missing columns for action '{a}' when merging model{k}: \"\n                        f\"{base} or {col_new} not found in merged probs\"\n                    )\n\n                probs_merged[base] = (\n                    probs_merged[base] * n_models + probs_merged[col_new]\n                ) / (n_models + 1)\n\n            # Drop the temporary per-model columns for this k\n            drop_cols_k = [f\"action_{a}_m{k}\" for a in ACTIONS]\n            probs_merged = probs_merged.drop(columns=drop_cols_k)\n\n            # Replace ensemble with updated merged/averaged version\n            probs_ens = probs_merged\n            n_models += 1\n\n            gc.collect()\n\n        # At this point probs_ens has equal-weighted average over (base + ENSEMBLE_N models)\n        if len(probs_ens) != len(probs):\n            log(\n                f\"{tag} WARNING: final merged rows {len(probs_ens)} != original {len(probs)} \"\n                f\"(some rows may be missing in one or more ensemble model files)\"\n            )\n\n        probs = probs_ens\n        del probs_ens\n        gc.collect()\n        #############################\n\n        V = probs.video_id.unique()\n        probs['lab_id'] = probs['video_id'].map(V2L).fillna(\"unknown\")\n        labb = probs.lab_id.values[0]\n\n        ### INSERT FILTER HERE ###\n        lab_key = str(labb)\n        allowed_actions = TRAIN_LAB_ACTIONS.get(lab_key)\n\n        if allowed_actions is not None and len(probs):\n            # Columns for actions that are allowed for THIS lab according to TRAIN\n            allowed_action_cols = {f\"action_{a}\" for a in allowed_actions}\n\n            # All action_* columns present\n            all_action_cols = [c for c in probs.columns if c.startswith(\"action_\")]\n\n            # Disallowed according to train: zero them out so they can't win argmax\n            disallowed_cols = [c for c in all_action_cols if c not in allowed_action_cols]\n            if disallowed_cols:\n                probs.loc[:, disallowed_cols] = 0.0\n                log(f\"{tag} lab={lab_key}: zeroed {len(disallowed_cols)} train-disallowed actions\")\n        else:\n            log(f\"{tag} lab={lab_key}: no TRAIN_LAB_ACTIONS entry (no extra filter)\")\n        ### END FILTER ###\n\n        ### BEGIN FILTER ###\n        CC = [f\"action_{a}\" for a in SELF_ACTIONS]\n        probs.loc[probs.agent_id != probs.target_id, CC] = 0.0\n\n        CC = [f\"action_{a}\" for a in PAIR_ACTIONS]\n        probs.loc[probs.agent_id == probs.target_id, CC] = 0.0\n        ### END FILTER ###\n        \n        try:\n            files = [f\"/tmp/xgb_preds/p{k}.pqt\" for k in V]\n            xgb = pd.read_parquet(files)\n            probs2 = probs.merge(xgb,on=['video_id','agent_id','target_id','frame'],how='left').fillna(0)\n            RMV = []\n            for a in ACTIONS:\n                n0 = f\"action_{a}\"\n                n1 = f\"action_{a}_x\"\n                n2 = f\"action_{a}_y\"\n                probs2[n0] = (1-XGB_WGT)*probs2[n1] + XGB_WGT*probs2[n2]\n                RMV.append(n1)\n                RMV.append(n2)\n            probs2 = probs2.drop(RMV,axis=1)\n            probs2 = probs2[COLS]\n        except:\n            log(f\"[XGB PREDS] no preds for {V}\")\n            probs2 = probs.copy()\n            for a in ACTIONS:\n                n0 = f\"action_{a}\"\n                probs2[n0] = (1-XGB_WGT)*probs2[n0] + XGB_WGT*0.2\n\n        # Smooth \n        smooth_probs_inplace(probs2, ACTIONS, win=SMOOTH_WIN)\n        log(f\"{tag} smooth done\")\n\n        # Mask by active set\n        if active_map_str and len(probs2):\n            probs2 = mask_probs_numpy_rle(probs2, ACTIONS, active_map_str, copy=False)\n            log(f\"{tag} mask done\")\n\n        # Intervals\n        thr_map = thresholds if thresholds else {a: DEFAULT_THR for a in ACTIONS}\n        sub = probs_to_nonoverlapping_intervals(probs2, ACTIONS, min_len=0, \n                                                max_gap=7, lab=labb, tie_config=TIE_CONFIG_V2)\n        del probs, probs2\n        gc.collect()\n\n        wrote = 0\n        if len(sub):\n            sub = sub[[\"video_id\",\"agent_id\",\"target_id\",\"action\",\"start_frame\",\"stop_frame\"]].copy()\n            sub = sub.pipe(lambda df: df.assign(\n                video_id=pd.to_numeric(df[\"video_id\"], errors=\"coerce\").fillna(-1).astype(int),\n                start_frame=pd.to_numeric(df[\"start_frame\"], errors=\"coerce\").fillna(0).astype(int),\n                stop_frame=pd.to_numeric(df[\"stop_frame\"], errors=\"coerce\").fillna(0).astype(int),\n                agent_id=pd.to_numeric(df[\"agent_id\"], errors=\"coerce\").fillna(0).astype(int),\n                target_id=pd.to_numeric(df[\"target_id\"], errors=\"coerce\").fillna(0).astype(int),\n                action=df[\"action\"].astype(str),\n            ))\n            sub = sub[sub[\"stop_frame\"] >= sub[\"start_frame\"]]\n            sub = sub[sub[\"action\"].isin(ACTIONS)]\n            # schema clean (writes \"self\" when agent==target)\n            def _format_agent_mouse(n: int) -> str: return f\"mouse{int(n)}\"\n            def _format_target_mouse(ai: int, ti: int) -> str: return \"self\" if int(ai)==int(ti) else f\"mouse{int(ti)}\"\n            sub[\"agent_id\"]  = sub[\"agent_id\"].apply(_format_agent_mouse)\n            sub[\"target_id\"] = [ _format_target_mouse(ai, ti) for ai,ti in zip(sub[\"agent_id\"].str.replace(\"mouse\",\"\").astype(int),\n                                                                               sub[\"target_id\"].astype(int)) ]\n            sub = sub[[\"video_id\",\"agent_id\",\"target_id\",\"action\",\"start_frame\",\"stop_frame\"]].drop_duplicates()\n            if len(sub):\n                sub = sub.reset_index(drop=True)\n                sub.insert(0, \"row_id\", range(row_counter, row_counter + len(sub)))\n                row_counter += len(sub)\n                sub.to_csv(OUT_PATH, mode=\"a\", header=False, index=False)\n                wrote = len(sub)\n        del sub\n        gc.collect()\n        if torch.cuda.is_available(): torch.cuda.empty_cache()\n\n        dt = time.time() - t_batch\n        log(f\"{tag} DONE in {dt:.1f}s | wrote_rows={wrote} | preds_total_rows_so_far={total_pred_rows} | { _gpu_mem() }\")\n\n    dt_all = time.time() - t0_all\n    log(f\"[done] wrote {OUT_PATH} | total_pred_rows={total_pred_rows} | total_time={dt_all:.1f}s | { _gpu_mem() }\")\n\nif __name__ == \"__main__\":\n    # Optional: improves conv/BN autotuning for fixed shapes\n    torch.backends.cudnn.benchmark = True\n    main()\n# ==================== END TEST-ONLY INFERENCE (LOW-RAM VERSION) ====================","metadata":{"collapsed":true,"execution":{"iopub.execute_input":"2025-12-15T13:59:23.030932Z","iopub.status.busy":"2025-12-15T13:59:23.030738Z","iopub.status.idle":"2025-12-15T14:00:00.471998Z","shell.execute_reply":"2025-12-15T14:00:00.471434Z"},"jupyter":{"outputs_hidden":true,"source_hidden":true},"papermill":{"duration":37.465061,"end_time":"2025-12-15T14:00:00.473197","exception":false,"start_time":"2025-12-15T13:59:23.008136","status":"completed"},"tags":[],"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\npd.read_csv(\"submission.csv\").sample(10)","metadata":{"execution":{"iopub.execute_input":"2025-12-15T14:00:00.535118Z","iopub.status.busy":"2025-12-15T14:00:00.534886Z","iopub.status.idle":"2025-12-15T14:00:00.546206Z","shell.execute_reply":"2025-12-15T14:00:00.54563Z"},"papermill":{"duration":0.043043,"end_time":"2025-12-15T14:00:00.547209","exception":false,"start_time":"2025-12-15T14:00:00.504166","status":"completed"},"tags":[]},"outputs":[],"execution_count":null}]}