{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10.17","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"tpu1vmV38","dataSources":[{"sourceId":105399,"databundleVersionId":12733338,"sourceType":"competition"}],"dockerImageVersionId":31042,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"%%capture\n!pip install -U xgboost\n!pip install -U polars\n!pip install -U optuna\n!pip install -U catboost\n!pip install -U lightgbm","metadata":{"_uuid":"c79029c5-49b2-44d8-8e67-9ed973f562ff","_cell_guid":"4aa06a0a-9a33-4d02-a78b-7780dc7081c7","trusted":true,"execution":{"iopub.status.busy":"2025-08-11T10:19:50.104126Z","iopub.execute_input":"2025-08-11T10:19:50.104513Z","iopub.status.idle":"2025-08-11T10:20:32.613203Z","shell.execute_reply.started":"2025-08-11T10:19:50.104483Z","shell.execute_reply":"2025-08-11T10:20:32.606461Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import polars as pl\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport time\nimport xgboost as xgb\nimport catboost\nimport lightgbm as lgb\nimport optuna\n\nRANDOM_STATE = 42\nnp.random.seed(RANDOM_STATE)","metadata":{"_uuid":"bdc84fc3-7c4c-47c1-82c6-8ceae3a83dd2","_cell_guid":"9fa8ceda-3412-4265-b1a5-e44a5711b17f","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-08-11T10:20:32.615177Z","iopub.execute_input":"2025-08-11T10:20:32.615418Z","iopub.status.idle":"2025-08-11T10:20:32.625363Z","shell.execute_reply.started":"2025-08-11T10:20:32.615391Z","shell.execute_reply":"2025-08-11T10:20:32.621196Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load data\ntrain = pl.read_parquet('/kaggle/input/aeroclub-recsys-2025/train.parquet').drop('__index_level_0__')\ntest = pl.read_parquet('/kaggle/input/aeroclub-recsys-2025/test.parquet').drop('__index_level_0__').with_columns(pl.lit(0, dtype=pl.Int64).alias(\"selected\"))\n\ndata_raw = pl.concat((train, test))","metadata":{"_uuid":"8cd97dac-cf5b-4c52-8bb1-f5559134339c","_cell_guid":"d83797d3-9898-4c98-a617-792950ac0adf","trusted":true,"execution":{"iopub.status.busy":"2025-08-11T10:20:32.626475Z","iopub.execute_input":"2025-08-11T10:20:32.626688Z","iopub.status.idle":"2025-08-11T10:20:36.488250Z","shell.execute_reply.started":"2025-08-11T10:20:32.626667Z","shell.execute_reply":"2025-08-11T10:20:36.482966Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Helpers","metadata":{"_uuid":"e6824382-eb9d-4c17-8ac7-dd656ee5b4b6","_cell_guid":"13df1e11-474b-4d7a-8da2-04c8c7e4c796","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"def hitrate_at_3(y_true, y_pred, groups):\n    df = pl.DataFrame({\n        'group': groups,\n        'pred': y_pred,\n        'true': y_true\n    })\n    \n    return (\n        df.filter(pl.col(\"group\").count().over(\"group\") > 10)\n        .sort([\"group\", \"pred\"], descending=[False, True])\n        .group_by(\"group\", maintain_order=True)\n        .head(3)\n        .group_by(\"group\")\n        .agg(pl.col(\"true\").max())\n        .select(pl.col(\"true\").mean())\n        .item()\n    )","metadata":{"_uuid":"1353ecee-4db7-4de8-a3eb-9fc33dc7040a","_cell_guid":"96d43bf4-2cbf-4fa4-bba2-72fb53e680ee","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-08-11T10:20:36.489292Z","iopub.execute_input":"2025-08-11T10:20:36.489535Z","iopub.status.idle":"2025-08-11T10:20:36.499801Z","shell.execute_reply.started":"2025-08-11T10:20:36.489511Z","shell.execute_reply":"2025-08-11T10:20:36.496347Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Feature Engineering","metadata":{"_uuid":"1fc1a4d6-5976-4b0e-9e28-4c8ff18e048a","_cell_guid":"d97a3b9b-bbb1-4f54-8889-959bb2c6f8ec","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"df = data_raw.clone()\n\n# More efficient duration to minutes converter\ndef dur_to_min(col):\n    # Extract days and time parts in one pass\n    days = col.str.extract(r\"^(\\d+)\\.\", 1).cast(pl.Int64).fill_null(0) * 1440\n    time_str = pl.when(col.str.contains(r\"^\\d+\\.\")).then(col.str.replace(r\"^\\d+\\.\", \"\")).otherwise(col)\n    hours = time_str.str.extract(r\"^(\\d+):\", 1).cast(pl.Int64).fill_null(0) * 60\n    minutes = time_str.str.extract(r\":(\\d+):\", 1).cast(pl.Int64).fill_null(0)\n    return (days + hours + minutes).fill_null(0)\n\n# Process duration columns\ndur_cols = [\"legs0_duration\", \"legs1_duration\"] + [f\"legs{l}_segments{s}_duration\" for l in (0, 1) for s in (0, 1)]\ndur_exprs = [dur_to_min(pl.col(c)).alias(c) for c in dur_cols if c in df.columns]\n\n# Apply duration transformations first\nif dur_exprs:\n    df = df.with_columns(dur_exprs)\n\n# Precompute marketing carrier columns check\nmc_cols = [f'legs{l}_segments{s}_marketingCarrier_code' for l in (0, 1) for s in range(4)]\nmc_exists = [col for col in mc_cols if col in df.columns]\n\n# Combine all initial transformations\ndf = df.with_columns([\n        # Price features\n        (pl.col(\"totalPrice\") / (pl.col(\"taxes\") + 1)).alias(\"price_per_tax\"),\n        (pl.col(\"taxes\") / (pl.col(\"totalPrice\") + 1)).alias(\"tax_rate\"),\n        pl.col(\"totalPrice\").log1p().alias(\"log_price\"),\n        \n        # Duration features\n        (pl.col(\"legs0_duration\").fill_null(0) + pl.col(\"legs1_duration\").fill_null(0)).alias(\"total_duration\"),\n        pl.when(pl.col(\"legs1_duration\").fill_null(0) > 0)\n            .then(pl.col(\"legs0_duration\") / (pl.col(\"legs1_duration\") + 1))\n            .otherwise(1.0).alias(\"duration_ratio\"),\n        \n        # Trip type\n        (pl.col(\"legs1_duration\").is_null() | \n         (pl.col(\"legs1_duration\") == 0) | \n         pl.col(\"legs1_segments0_departureFrom_airport_iata\").is_null()).cast(pl.Int32).alias(\"is_one_way\"),\n        \n        # Total segments count\n        (pl.sum_horizontal(pl.col(col).is_not_null().cast(pl.UInt8) for col in mc_exists) \n         if mc_exists else pl.lit(0)).alias(\"l0_seg\"),\n        \n        # FF features\n        (pl.col(\"frequentFlyer\").fill_null(\"\").str.count_matches(\"/\") + \n         (pl.col(\"frequentFlyer\").fill_null(\"\") != \"\").cast(pl.Int32)).alias(\"n_ff_programs\"),\n        \n        # Binary features\n        pl.col(\"corporateTariffCode\").is_not_null().cast(pl.Int32).alias(\"has_corporate_tariff\"),\n        (pl.col(\"pricingInfo_isAccessTP\") == 1).cast(pl.Int32).alias(\"has_access_tp\"),\n        \n        # Baggage & fees\n        (pl.col(\"legs0_segments0_baggageAllowance_quantity\").fill_null(0) + \n         pl.col(\"legs1_segments0_baggageAllowance_quantity\").fill_null(0)).alias(\"baggage_total\"),\n        (pl.col(\"miniRules0_monetaryAmount\").fill_null(0) + \n         pl.col(\"miniRules1_monetaryAmount\").fill_null(0)).alias(\"total_fees\"),\n        \n        # Routes & carriers\n        pl.col(\"searchRoute\").is_in([\"MOWLED/LEDMOW\", \"LEDMOW/MOWLED\", \"MOWLED\", \"LEDMOW\", \"MOWAER/AERMOW\"])\n            .cast(pl.Int32).alias(\"is_popular_route\"),\n        \n        # Cabin\n        pl.mean_horizontal([\"legs0_segments0_cabinClass\", \"legs1_segments0_cabinClass\"]).alias(\"avg_cabin_class\"),\n        (pl.col(\"legs0_segments0_cabinClass\").fill_null(0) - \n         pl.col(\"legs1_segments0_cabinClass\").fill_null(0)).alias(\"cabin_class_diff\"),\n])\n\n# Segment counts - more efficient\nseg_exprs = []\nfor leg in (0, 1):\n    seg_cols = [f\"legs{leg}_segments{s}_duration\" for s in range(4) if f\"legs{leg}_segments{s}_duration\" in df.columns]\n    if seg_cols:\n        seg_exprs.append(\n            pl.sum_horizontal(pl.col(c).is_not_null() for c in seg_cols)\n                .cast(pl.Int32).alias(f\"n_segments_leg{leg}\")\n        )\n    else:\n        seg_exprs.append(pl.lit(0).cast(pl.Int32).alias(f\"n_segments_leg{leg}\"))\n\n# Add segment-based features\n# First create segment counts\ndf = df.with_columns(seg_exprs)\n\n# Then use them for derived features\ndf = df.with_columns([\n    (pl.col(\"n_segments_leg0\") + pl.col(\"n_segments_leg1\")).alias(\"total_segments\"),\n    (pl.col(\"n_segments_leg0\") == 1).cast(pl.Int32).alias(\"is_direct_leg0\"),\n    pl.when(pl.col(\"is_one_way\") == 1).then(0)\n        .otherwise((pl.col(\"n_segments_leg1\") == 1).cast(pl.Int32)).alias(\"is_direct_leg1\"),\n])\n\n# More derived features\ndf = df.with_columns([\n    (pl.col(\"is_direct_leg0\") & pl.col(\"is_direct_leg1\")).cast(pl.Int32).alias(\"both_direct\"),\n    ((pl.col(\"isVip\") == 1) | (pl.col(\"n_ff_programs\") > 0)).cast(pl.Int32).alias(\"is_vip_freq\"),\n    (pl.col(\"baggage_total\") > 0).cast(pl.Int32).alias(\"has_baggage\"),\n    (pl.col(\"total_fees\") > 0).cast(pl.Int32).alias(\"has_fees\"),\n    (pl.col(\"total_fees\") / (pl.col(\"totalPrice\") + 1)).alias(\"fee_rate\"),\n    pl.col(\"Id\").count().over(\"ranker_id\").alias(\"group_size\"),\n])\n\n# Add major carrier flag if column exists\nif \"legs0_segments0_marketingCarrier_code\" in df.columns:\n    df = df.with_columns(\n        pl.col(\"legs0_segments0_marketingCarrier_code\").is_in([\"SU\", \"S7\", \"U6\"])\n            .cast(pl.Int32).alias(\"is_major_carrier\")\n    )\nelse:\n    df = df.with_columns(pl.lit(0).alias(\"is_major_carrier\"))\n\ndf = df.with_columns(pl.col(\"group_size\").log1p().alias(\"group_size_log\"))\n\n# Time features - batch process\ntime_exprs = []\nfor col in (\"legs0_departureAt\", \"legs0_arrivalAt\", \"legs1_departureAt\", \"legs1_arrivalAt\"):\n    if col in df.columns:\n        dt = pl.col(col).str.to_datetime(strict=False)\n        h = dt.dt.hour().fill_null(12)\n        time_exprs.extend([\n            h.alias(f\"{col}_hour\"),\n            dt.dt.weekday().fill_null(0).alias(f\"{col}_weekday\"),\n            (((h >= 6) & (h <= 9)) | ((h >= 17) & (h <= 20))).cast(pl.Int32).alias(f\"{col}_business_time\")\n        ])\nif time_exprs:\n    df = df.with_columns(time_exprs)\n\n# Batch rank computations - more efficient with single pass\n# First apply the columns that will be used for ranking\ndf = df.with_columns([\n    pl.col(\"group_size\").log1p().alias(\"group_size_log\"),\n])\n\n# Price and duration basic ranks\nrank_exprs = []\nfor col, alias in [(\"totalPrice\", \"price\"), (\"total_duration\", \"duration\")]:\n    rank_exprs.append(pl.col(col).rank().over(\"ranker_id\").alias(f\"{alias}_rank\"))\n\n# Price-specific features\nprice_exprs = [\n    (pl.col(\"totalPrice\").rank(\"average\").over(\"ranker_id\") / \n     pl.col(\"totalPrice\").count().over(\"ranker_id\")).alias(\"price_pct_rank\"),\n    (pl.col(\"totalPrice\") == pl.col(\"totalPrice\").min().over(\"ranker_id\")).cast(pl.Int32).alias(\"is_cheapest\"),\n    ((pl.col(\"totalPrice\") - pl.col(\"totalPrice\").median().over(\"ranker_id\")) / \n     (pl.col(\"totalPrice\").std().over(\"ranker_id\") + 1)).alias(\"price_from_median\"),\n    (pl.col(\"l0_seg\") == pl.col(\"l0_seg\").min().over(\"ranker_id\")).cast(pl.Int32).alias(\"is_min_segments\"),\n]\n\n# Apply initial ranks\ndf = df.with_columns(rank_exprs + price_exprs)\n\n# Cheapest direct - more efficient\ndirect_cheapest = (\n    df.filter(pl.col(\"is_direct_leg0\") == 1)\n    .group_by(\"ranker_id\")\n    .agg(pl.col(\"totalPrice\").min().alias(\"min_direct\"))\n)\n\ndf = df.join(direct_cheapest, on=\"ranker_id\", how=\"left\").with_columns(\n    ((pl.col(\"is_direct_leg0\") == 1) & \n     (pl.col(\"totalPrice\") == pl.col(\"min_direct\"))).cast(pl.Int32).fill_null(0).alias(\"is_direct_cheapest\")\n).drop(\"min_direct\")","metadata":{"_uuid":"98d8eede-ea65-47a2-860f-c7547b2f761a","_cell_guid":"50d34c94-36c7-44e7-8019-ffb0b1633255","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-08-11T10:20:36.502476Z","iopub.execute_input":"2025-08-11T10:20:36.502953Z","iopub.status.idle":"2025-08-11T10:21:39.383299Z","shell.execute_reply.started":"2025-08-11T10:20:36.502929Z","shell.execute_reply":"2025-08-11T10:21:39.378019Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"## NEW: Advanced Feature Engineering: Group-wise Interaction Features\nprint(\"Starting advanced group-wise feature engineering...\")\n\n# Define the main carrier column, handling potential missing columns\nmain_carrier_col = \"legs0_segments0_marketingCarrier_code\"\nif main_carrier_col not in df.columns:\n    df = df.with_columns(pl.lit(None, dtype=pl.String).alias(main_carrier_col))\n\n# Calculate group-wise statistics in a separate dataframe\ngroup_stats = df.group_by(\"ranker_id\").agg(\n    pl.col(\"totalPrice\").min().alias(\"min_price_group\"),\n    pl.col(\"totalPrice\").max().alias(\"max_price_group\"),\n    pl.col(\"totalPrice\").mean().alias(\"mean_price_group\"),\n    pl.col(\"total_duration\").min().alias(\"min_duration_group\"),\n    pl.col(main_carrier_col).mode().first().alias(\"mode_carrier_group\")\n)\n\n# Join the group stats back to the main dataframe\ndf = df.join(group_stats, on=\"ranker_id\", how=\"left\")\n\n# --- FIX IS HERE ---\n\n# Step 1: Create the initial set of interaction features\ndf = df.with_columns(\n    # Price interaction features\n    (pl.col(\"totalPrice\") - pl.col(\"min_price_group\")).alias(\"price_vs_min\"),\n    (pl.col(\"max_price_group\") - pl.col(\"totalPrice\")).alias(\"price_vs_max\"),\n    (pl.col(\"totalPrice\") / pl.col(\"mean_price_group\")).alias(\"price_vs_mean\"),\n    \n    # Duration interaction features\n    (pl.col(\"total_duration\") - pl.col(\"min_duration_group\")).alias(\"duration_vs_min\"),\n    \n    # Carrier interaction features\n    (pl.col(main_carrier_col) == pl.col(\"mode_carrier_group\")).cast(pl.Int32).alias(\"is_mode_carrier\")\n)\n\n# Step 2: Now that 'price_vs_min' exists, create the combined feature\ndf = df.with_columns(\n    (pl.col(\"price_vs_min\") / (pl.col(\"total_duration\") + 1)).alias(\"price_vs_min_per_duration\")\n)\n\n# --- END OF FIX ---\n\n# Drop the intermediate stats columns\ndf = df.drop([\"min_price_group\", \"max_price_group\", \"mean_price_group\", \"min_duration_group\", \"mode_carrier_group\"])\n\nprint(\"Advanced features created successfully.\")","metadata":{"_uuid":"abe88f4c-b47f-4958-902d-138b6cf5351f","_cell_guid":"4d0ae7f1-9aec-4e48-950c-b33ed6d899c5","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-08-11T10:21:39.385439Z","iopub.execute_input":"2025-08-11T10:21:39.385680Z","iopub.status.idle":"2025-08-11T10:21:40.981371Z","shell.execute_reply.started":"2025-08-11T10:21:39.385655Z","shell.execute_reply":"2025-08-11T10:21:40.976727Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Fill nulls\ndata = df.with_columns(\n    [pl.col(c).fill_null(0) for c in df.select(pl.selectors.numeric()).columns] +\n    [pl.col(c).fill_null(\"missing\") for c in df.select(pl.selectors.string()).columns]\n)","metadata":{"_uuid":"bb77cfd4-d82c-4b25-a6f0-dc252511812b","_cell_guid":"ca300ed1-737d-4e36-890a-7a10b045f535","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-08-11T10:21:40.984051Z","iopub.execute_input":"2025-08-11T10:21:40.984327Z","iopub.status.idle":"2025-08-11T10:21:44.661492Z","shell.execute_reply.started":"2025-08-11T10:21:40.984300Z","shell.execute_reply":"2025-08-11T10:21:44.656128Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Feature Selection","metadata":{"_uuid":"5ef9eeea-c1cd-4e0e-9406-70c0c4cdb30c","_cell_guid":"0662d676-370a-4df3-a225-a41f16855599","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"# Categorical features\ncat_features = [\n    'nationality', 'searchRoute', 'corporateTariffCode',\n    'bySelf', 'sex', 'companyID',\n    # Leg 0 segments 0-1\n    'legs0_segments0_aircraft_code', 'legs0_segments0_arrivalTo_airport_city_iata',\n    'legs0_segments0_arrivalTo_airport_iata', 'legs0_segments0_departureFrom_airport_iata',\n    'legs0_segments0_marketingCarrier_code', 'legs0_segments0_operatingCarrier_code',\n    'legs0_segments0_flightNumber',\n    'legs0_segments1_aircraft_code', 'legs0_segments1_arrivalTo_airport_city_iata',\n    'legs0_segments1_arrivalTo_airport_iata', 'legs0_segments1_departureFrom_airport_iata',\n    'legs0_segments1_marketingCarrier_code', 'legs0_segments1_operatingCarrier_code',\n    'legs0_segments1_flightNumber',\n    # Leg 1 segments 0-1\n    'legs1_segments0_aircraft_code', 'legs1_segments0_arrivalTo_airport_city_iata',\n    'legs1_segments0_arrivalTo_airport_iata', 'legs1_segments0_departureFrom_airport_iata',\n    'legs1_segments0_marketingCarrier_code', 'legs1_segments0_operatingCarrier_code',\n    'legs1_segments0_flightNumber',\n    'legs1_segments1_aircraft_code', 'legs1_segments1_arrivalTo_airport_city_iata',\n    'legs1_segments1_arrivalTo_airport_iata', 'legs1_segments1_departureFrom_airport_iata',\n    'legs1_segments1_marketingCarrier_code', 'legs1_segments1_operatingCarrier_code',\n    'legs1_segments1_flightNumber',\n]\n\n# Columns to exclude (uninformative or problematic)\nexclude_cols = [\n    'Id', 'ranker_id', 'selected', 'profileId', 'requestDate',\n    'legs0_departureAt', 'legs0_arrivalAt', 'legs1_departureAt', 'legs1_arrivalAt',\n    'miniRules0_percentage', 'miniRules1_percentage',  # >90% missing\n    'frequentFlyer',  # Already processed\n    # Exclude constant columns\n    'pricingInfo_passengerCount'\n]\n\n\n# Exclude segment 2-3 columns (>98% missing)\nfor leg in [0, 1]:\n    for seg in [2, 3]:\n        for suffix in ['aircraft_code', 'arrivalTo_airport_city_iata', 'arrivalTo_airport_iata',\n                      'baggageAllowance_quantity', 'baggageAllowance_weightMeasurementType',\n                      'cabinClass', 'departureFrom_airport_iata', 'duration', 'flightNumber',\n                      'marketingCarrier_code', 'operatingCarrier_code', 'seatsAvailable']:\n            exclude_cols.append(f'legs{leg}_segments{seg}_{suffix}')\n\nfeature_cols = [col for col in data.columns if col not in exclude_cols]\ncat_features_final = [col for col in cat_features if col in feature_cols]\n\nprint(f\"Using {len(feature_cols)} features ({len(cat_features_final)} categorical)\")\n\nX = data.select(feature_cols)\ny = data.select('selected')\ngroups = data.select('ranker_id')","metadata":{"_uuid":"582b2021-cef1-4203-8523-f5a27dda39a6","_cell_guid":"8c24560e-8c5e-4985-bf6f-b142a1d9db3e","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-08-11T10:21:44.663979Z","iopub.execute_input":"2025-08-11T10:21:44.664262Z","iopub.status.idle":"2025-08-11T10:21:44.679646Z","shell.execute_reply.started":"2025-08-11T10:21:44.664236Z","shell.execute_reply":"2025-08-11T10:21:44.676441Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Model Training: Two-Stage Re-ranking","metadata":{"_uuid":"ca2f0045-cad3-4f6c-9d8c-2789ecaacd01","_cell_guid":"e0ce949e-ba1c-402e-bfde-438ed514232f","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"## NEW: Prepare data for modeling\n# Encode categorical features for compatibility with both LGBM and XGB\nX_encoded = X.with_columns([(pl.col(c).rank(\"dense\") - 1).fill_null(-1).cast(pl.Int32) for c in cat_features_final])\n\n# Define train/val/test splits\nn1 = 16487352 # split train to train and val (10%) in time\nn2 = train.height\n\n# Full datasets\nX_tr, X_va, X_te = X_encoded[:n1], X_encoded[n1:n2], X_encoded[n2:]\ny_tr, y_va, y_te = y[:n1], y[n1:n2], y[n2:]\ngroups_tr, groups_va, groups_te = groups[:n1], groups[n1:n2], groups[n2:]\n\n# Get group sizes for the full datasets\ngroup_sizes_tr = groups_tr.group_by('ranker_id').agg(pl.len()).sort('ranker_id')['len'].to_numpy()\ngroup_sizes_va = groups_va.group_by('ranker_id').agg(pl.len()).sort('ranker_id')['len'].to_numpy()","metadata":{"_uuid":"e1e11949-58b9-4b5b-8913-ca595eb85c8e","_cell_guid":"97cb77bf-964a-405f-aea2-1fc03b3ba6b7","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-08-11T10:21:44.682358Z","iopub.execute_input":"2025-08-11T10:21:44.682657Z","iopub.status.idle":"2025-08-11T10:21:54.022397Z","shell.execute_reply.started":"2025-08-11T10:21:44.682634Z","shell.execute_reply":"2025-08-11T10:21:54.014555Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Stage 1: First-Pass Ranking with LightGBM\nThe goal of this model is to quickly identify a smaller set of promising candidates from the full list of options.","metadata":{"_uuid":"e82e021f-682e-424d-9f5f-801041dde7a0","_cell_guid":"2cba9694-cfc2-4690-960b-60d821d2b369","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"## NEW: Train Stage 1 Model\nprint(\"--- Training Stage 1: LightGBM First-Pass Ranker ---\")\n\n# Create LightGBM datasets\nlgb_train_s1 = lgb.Dataset(\n    data=X_tr, \n    label=y_tr.to_numpy().flatten(), \n    group=group_sizes_tr,\n    feature_name=feature_cols,\n    free_raw_data=False\n)\n\nlgb_val_s1 = lgb.Dataset(\n    data=X_va, \n    label=y_va.to_numpy().flatten(), \n    group=group_sizes_va,\n    feature_name=feature_cols,\n    reference=lgb_train_s1,\n    free_raw_data=False\n)\n\n# Use a relatively simple and fast LightGBM configuration\ns1_lgb_params = {\n    'objective': 'lambdarank', 'metric': 'ndcg', 'boosting_type': 'gbdt','eval_at': [3],\n    'num_leaves': 40, 'learning_rate': 0.1, 'n_estimators': 500,\n    'feature_fraction': 0.8, 'bagging_fraction': 0.8, 'bagging_freq': 1,\n    'n_jobs': -1, 'random_state': RANDOM_STATE, 'label_gain': [0, 1]\n}\n\nlgb_model_s1 = lgb.train(\n    s1_lgb_params,\n    lgb_train_s1,\n    valid_sets=[lgb_train_s1, lgb_val_s1],\n    callbacks=[lgb.early_stopping(50), lgb.log_evaluation(100)]\n)","metadata":{"_uuid":"c0255bd1-2fce-4283-99d0-e587f6ab7647","_cell_guid":"d4dba269-0007-4b33-8b05-2d090790d5ca","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-08-11T10:21:54.024950Z","iopub.execute_input":"2025-08-11T10:21:54.025392Z","iopub.status.idle":"2025-08-11T10:26:08.778468Z","shell.execute_reply.started":"2025-08-11T10:21:54.025354Z","shell.execute_reply":"2025-08-11T10:26:08.772903Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Prepare Data for Stage 2\nWe now use the Stage 1 model to predict on the training and validation sets. We select the top `k` candidates from each group to create a new, smaller dataset for the Stage 2 model.","metadata":{"_uuid":"e52bf74a-8df8-444e-a992-a23463232f68","_cell_guid":"a7eb9a92-d6c6-4be9-a5b7-5da2a9b86a8b","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"## NEW: Filter data for Stage 2\nTOP_K = 15 # Number of candidates to select for the re-ranking stage\nprint(f\"\\n--- Filtering for Stage 2: Selecting Top {TOP_K} Candidates ---\")\n\ndef filter_top_k(data_X, data_y, data_groups, model, k):\n    \"\"\"Uses a trained model to predict and filter for the top k candidates per group.\"\"\"\n    # Get predictions\n    preds = model.predict(data_X)\n    \n    # Create a temporary dataframe with predictions\n    df_preds = data_groups.with_columns(\n        pl.Series(\"preds\", preds),\n        pl.Series(\"true\", data_y.to_series()),\n        pl.Series(\"row_nr\", np.arange(len(data_X))) # Original index\n    )\n    \n    # Get the original row numbers of the top k candidates for each group\n    top_k_indices = (\n        df_preds\n        .sort(\"preds\", descending=True)\n        .group_by(\"ranker_id\")\n        .head(k)\n        .get_column(\"row_nr\")\n        .to_numpy()\n    )\n    \n    # Filter the original dataframes using these indices\n    X_rerank = data_X[top_k_indices]\n    y_rerank = data_y[top_k_indices]\n    groups_rerank = data_groups[top_k_indices]\n    \n    return X_rerank, y_rerank, groups_rerank\n\n# Create the re-ranking training set from the original training data\nX_tr_rerank, y_tr_rerank, groups_tr_rerank = filter_top_k(X_tr, y_tr, groups_tr, lgb_model_s1, TOP_K)\n\n# Create the re-ranking validation set from the original validation data\nX_va_rerank, y_va_rerank, groups_va_rerank = filter_top_k(X_va, y_va, groups_va, lgb_model_s1, TOP_K)\n\nprint(f\"Original training size: {len(X_tr)}, Re-ranking training size: {len(X_tr_rerank)}\")\nprint(f\"Original validation size: {len(X_va)}, Re-ranking validation size: {len(X_va_rerank)}\")","metadata":{"_uuid":"99871c73-07cc-4a97-83c9-30753d4cec73","_cell_guid":"18b56eda-bc8a-415f-8eda-a4ecdff9240d","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-08-11T10:26:08.779339Z","iopub.execute_input":"2025-08-11T10:26:08.779607Z","iopub.status.idle":"2025-08-11T10:27:35.963487Z","shell.execute_reply.started":"2025-08-11T10:26:08.779578Z","shell.execute_reply":"2025-08-11T10:27:35.958795Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Stage 2: Re-ranking with XGBoost\nNow we train a more powerful XGBoost model on the smaller, higher-quality dataset. This model can learn the finer distinctions between the best options.","metadata":{"_uuid":"39500463-a477-4ea7-ba63-08239647ed33","_cell_guid":"39468e37-eb6c-440c-8c01-56852a83048c","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"## NEW: Train Stage 2 Model\nprint(\"\\n--- Training Stage 2: XGBoost Re-Ranker ---\")\n\n# Get group sizes for the new re-ranking datasets\ngroup_sizes_tr_rerank = groups_tr_rerank.group_by('ranker_id').agg(pl.len()).sort('ranker_id')['len'].to_numpy()\ngroup_sizes_va_rerank = groups_va_rerank.group_by('ranker_id').agg(pl.len()).sort('ranker_id')['len'].to_numpy()\n\n# Create XGBoost DMatrix objects for the re-ranking data\ndtrain_rerank = xgb.DMatrix(X_tr_rerank, label=y_tr_rerank, group=group_sizes_tr_rerank, feature_names=X_tr_rerank.columns)\ndval_rerank   = xgb.DMatrix(X_va_rerank, label=y_va_rerank, group=group_sizes_va_rerank, feature_names=X_va_rerank.columns)\n\n# Use your optimized XGBoost parameters\ns2_xgb_params = {\n    'objective': 'rank:pairwise', 'eval_metric': 'ndcg@3', \n    'max_depth': 8, 'min_child_weight': 14, 'subsample': 0.9, \n    'colsample_bytree': 1.0, 'lambda': 3.5330891736457763 , \n    'learning_rate': 0.0521879929228514 ,\n    'seed': RANDOM_STATE, 'n_jobs': -1\n}\n\nprint(\"\\nTraining final XGBoost re-ranking model...\")\nxgb_model_s2 = xgb.train(\n    s2_xgb_params,\n    dtrain_rerank,\n    num_boost_round=2000, # Can increase boosting rounds as the dataset is smaller\n    evals=[(dtrain_rerank, 'train'), (dval_rerank, 'val')],\n    early_stopping_rounds=150, # More patience for the final model\n    verbose_eval=50\n)","metadata":{"_uuid":"132950c7-a4b3-4558-baad-23f25c8aaa1f","_cell_guid":"ac70ae44-1a3f-46e1-8507-32082b66395b","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-08-11T10:27:35.966225Z","iopub.execute_input":"2025-08-11T10:27:35.966461Z","iopub.status.idle":"2025-08-11T10:29:05.394970Z","shell.execute_reply.started":"2025-08-11T10:27:35.966437Z","shell.execute_reply":"2025-08-11T10:29:05.390863Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Final Evaluation","metadata":{"_uuid":"26374ccd-c3b0-4526-8aa4-4d88e29a7dc4","_cell_guid":"e9bee7b6-55e7-4326-9c7b-eb99372b5ded","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"## NEW: Evaluate the Stage 2 model\nprint(\"\\n--- Evaluating Re-ranking Model on Validation Set ---\")\n\n# Predict with the Stage 2 model on the re-ranking validation set\nrerank_va_preds = xgb_model_s2.predict(dval_rerank)\n\n# Calculate HitRate@3\nrerank_hr3 = hitrate_at_3(\n    y_va_rerank['selected'], \n    rerank_va_preds, \n    groups_va_rerank['ranker_id']\n)\n\nprint(\"-\" * 30)\nprint(f\"Stage 2 XGBoost Re-ranker HitRate@3: {rerank_hr3:.4f}\")\nprint(\"-\" * 30)","metadata":{"_uuid":"0d325f09-52ce-4154-8854-d5ffe8e47076","_cell_guid":"647fb43e-f6c2-4504-a0fb-71e083fc9114","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-08-11T10:29:05.396905Z","iopub.execute_input":"2025-08-11T10:29:05.397171Z","iopub.status.idle":"2025-08-11T10:29:05.531724Z","shell.execute_reply.started":"2025-08-11T10:29:05.397133Z","shell.execute_reply":"2025-08-11T10:29:05.527221Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Submission\nThe submission process now also follows the two-stage approach:\n1. Use the Stage 1 model to select the Top K candidates from the test set.\n2. Use the Stage 2 model to re-rank only these candidates and generate the final submission.","metadata":{"_uuid":"beebbe37-2cb6-43b0-99a7-c63114907ba2","_cell_guid":"ad84cfaa-0fb9-4dc5-8281-62b1de8c03b6","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}}},{"cell_type":"code","source":"## FINAL: Generate submission using the two-stage pipeline (Rank Fix)\nprint(\"\\n--- Generating Submission File using Two-Stage Pipeline ---\")\n\n# Stage 1: Predict on the full test set to find top K candidates\nprint(f\"Stage 1: Filtering test set to Top {TOP_K} candidates...\")\ns1_test_preds = lgb_model_s1.predict(X_te)\n\n# Create dataframe with predictions and original row numbers\ndf_test_preds = groups_te.with_columns(\n    pl.Series(\"preds\", s1_test_preds),\n    pl.Series(\"row_nr\", np.arange(len(X_te)))\n)\n\n# Get top K indices per ranker_id\ntop_k_test_indices = (\n    df_test_preds\n    .sort(\"preds\", descending=True)\n    .group_by(\"ranker_id\")\n    .head(TOP_K)\n    .get_column(\"row_nr\")\n    .to_numpy()\n)\n\n# Filter test set to top K only\nX_te_rerank = X_te[top_k_test_indices]\ntest_rerank = test[top_k_test_indices]\n\n# Stage 2: Re-rank candidates using XGBoost\nprint(\"Stage 2: Re-ranking candidates with XGBoost model...\")\ndtest_rerank = xgb.DMatrix(X_te_rerank, feature_names=X_te_rerank.columns)\ns2_test_preds = xgb_model_s2.predict(dtest_rerank)\n\n# Create ranked dataframe\nsubmission_df = test_rerank.select(['Id', 'ranker_id']).with_columns(\n    pl.Series('final_score', s2_test_preds)\n).with_columns(\n    pl.col('final_score')\n      .rank(method='ordinal', descending=True)\n      .over('ranker_id')\n      .cast(pl.Int32)\n      .alias('selected')\n).select(['Id', 'ranker_id', 'selected'])\n\n# Merge with full test set and fill non-top-K with rank 99\nfull_submission = (\n    test.select(['Id', 'ranker_id'])\n    .join(submission_df, on=['Id', 'ranker_id'], how='left')\n    .with_columns(\n        pl.col('selected').fill_null(99)\n    )\n)\n\n# Save CSV\nfull_submission.write_csv('submission.csv')\n\nprint(\"\\n✅ Submission file 'submission.csv' created successfully.\")\nprint(full_submission.sort(\"ranker_id\").head(20))\n","metadata":{"_uuid":"2e2643d4-057a-466e-a93e-5704fd2e9ab9","_cell_guid":"04faa3a3-ab17-4744-9159-125acba90102","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false},"execution":{"iopub.status.busy":"2025-08-11T10:29:05.534055Z","iopub.execute_input":"2025-08-11T10:29:05.534585Z","iopub.status.idle":"2025-08-11T10:29:40.934737Z","shell.execute_reply.started":"2025-08-11T10:29:05.534558Z","shell.execute_reply":"2025-08-11T10:29:40.928533Z"}},"outputs":[],"execution_count":null}]}