Harness context codex item completed
{
"type": "codex_item_completed",
"item": {
"type": "CommandExecution",
"id": "exec-21e2ff5d-6c78-456c-af62-edfa2366f842",
"process_id": "61724",
"command": [
"/bin/sh",
"-lc",
"sed -n '760,1225p' /app/solution/train_current_model.py\nsed -n '1225,1790p' /app/solution/train_current_model.py\nsed -n '1790,2095p' /app/solution/train_current_model.py"
],
"cwd": "file:///app",
"parsed_cmd": [
{
"type": "read",
"cmd": "sed -n '760,1225p' /app/solution/train_current_model.py",
"name": "train_current_model.py",
"path": "/app/solution/train_current_model.py"
},
{
"type": "read",
"cmd": "sed -n '1225,1790p' /app/solution/train_current_model.py",
"name": "train_current_model.py",
"path": "/app/solution/train_current_model.py"
},
{
"type": "read",
"cmd": "sed -n '1790,2095p' /app/solution/train_current_model.py",
"name": "train_current_model.py",
"path": "/app/solution/train_current_model.py"
}
],
"source": "unified_exec_startup",
"status": "completed",
"stdout": " )\n\n\ndef event_sequence(\n model: base.SequentialTwoTower,\n items: torch.Tensor,\n features: torch.Tensor,\n hidden: torch.Tensor,\n) -> tuple[torch.Tensor, torch.Tensor]:\n inputs = torch.cat([model.context_embedding(items), features], dim=-1)\n projected = model.event_projection(inputs)\n return model.sequence_encoder(projected, hidden)\n\n\ndef pool_explicit_states(\n model: base.SequentialTwoTower,\n prepared: SequenceData,\n state_ids: np.ndarray,\n device: torch.device,\n) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n unique_states, inverse = np.unique(state_ids, return_inverse=True)\n liked_items = torch.from_numpy(\n prepared.liked_states[unique_states].astype(np.int64)\n ).to(device)\n disliked_items = torch.from_numpy(\n prepared.disliked_states[unique_states].astype(np.int64)\n ).to(device)\n liked_vectors = model._mean_context_channel(liked_items)\n disliked_vectors = model._mean_context_channel(disliked_items)\n inverse_rows = torch.from_numpy(inverse.astype(np.int64)).to(device)\n counts = prepared.state_counts[state_ids]\n return (\n liked_vectors[inverse_rows],\n disliked_vectors[inverse_rows],\n torch.from_numpy(counts.astype(np.float32)).to(device),\n )\n\n\ndef user_queries(\n model: base.SequentialTwoTower,\n sequence_vectors: torch.Tensor,\n liked_vectors: torch.Tensor,\n disliked_vectors: torch.Tensor,\n user_features: torch.Tensor,\n) -> torch.Tensor:\n return model.user_tower(\n torch.cat(\n [sequence_vectors, liked_vectors, disliked_vectors, user_features],\n dim=1,\n )\n )\n\n\ndef candidate_vectors(\n model: base.SequentialTwoTower,\n positives: torch.Tensor,\n negatives: torch.Tensor,\n) -> tuple[torch.Tensor, torch.Tensor]:\n all_items = torch.cat([positives, negatives.reshape(-1)])\n unique_items, inverse = torch.unique(all_items, return_inverse=True)\n encoded = model.encode_items(unique_items)\n positive_vectors = encoded[inverse[: len(positives)]]\n negative_vectors = encoded[inverse[len(positives) :]].reshape(\n len(positives), negatives.shape[1], -1\n )\n return positive_vectors, negative_vectors\n\n\ndef history_windows(\n sequence: np.ndarray,\n start: int,\n count: int,\n width: int,\n padding_item: int,\n) -> np.ndarray:\n base_start = max(0, start - width + 1)\n segment = sequence[base_start : start + count]\n padded = np.concatenate(\n [np.full(width - 1, padding_item, dtype=np.uint16), segment]\n )\n windows = np.lib.stride_tricks.sliding_window_view(padded, width)\n offset = start - base_start\n return windows[offset : offset + count].copy()\n\n\ndef corrected_negative_loss(\n logits: torch.Tensor,\n correction: torch.Tensor,\n config: Config,\n) -> torch.Tensor:\n raw = F.softplus(logits)\n total = config.uniform_negatives + config.inbatch_negatives\n result = torch.zeros(len(logits), device=logits.device)\n if config.uniform_negatives:\n result += config.uniform_negatives * raw[\n :, : config.uniform_negatives\n ].mean(dim=1)\n if config.inbatch_negatives:\n selected = slice(config.uniform_negatives, total)\n weights = correction[:, selected]\n result += config.inbatch_negatives * (\n (raw[:, selected] * weights).sum(dim=1)\n / weights.sum(dim=1).clamp_min(1.0e-6)\n )\n return result / total\n\n\ndef retained_target_window_starts(\n total_targets: np.ndarray,\n validation_counts: np.ndarray,\n training_counts: np.ndarray,\n) -> np.ndarray:\n \"\"\"Return the discarded prefix before the retained chronological suffix.\"\"\"\n training_region_ends = total_targets - validation_counts\n window_starts = training_region_ends - training_counts\n if (\n np.any(window_starts < 0)\n or np.any(training_counts < 0)\n or np.any(validation_counts < 0)\n ):\n raise RuntimeError(\"chronological target partition produced negative counts\")\n if not np.array_equal(\n window_starts + training_counts + validation_counts, total_targets\n ):\n raise RuntimeError(\"chronological target window is not a final contiguous suffix\")\n return window_starts\n\n\ndef chronological_target_partition(\n prepared: SequenceData,\n validation_fraction: float,\n max_train_targets_per_user: int,\n max_validation_targets_per_user: int,\n) -> TargetPartition:\n \"\"\"Select a recent contiguous window with a final out-of-time BCE suffix.\"\"\"\n totals = np.maximum(prepared.lengths - 1, 0).astype(np.int64)\n validation = np.floor(totals * validation_fraction).astype(np.int64)\n validation[(totals >= 2) & (validation == 0)] = 1\n if max_validation_targets_per_user:\n validation = np.minimum(validation, max_validation_targets_per_user)\n\n history_training = totals - validation\n training = history_training.copy()\n if max_train_targets_per_user:\n training = np.minimum(training, max_train_targets_per_user)\n window_starts = retained_target_window_starts(totals, validation, training)\n return TargetPartition(\n window_starts=window_starts,\n training_counts=training,\n validation_counts=validation,\n history_training_counts=history_training,\n )\n\n\ndef load_organic_like_outcomes(\n public: Path, stage: str, item_ids: np.ndarray\n) -> dict[int, dict[int, np.ndarray]]:\n \"\"\"Load mapped organic-like timestamps from public fitting history.\"\"\"\n prefixes = (\"train\",) if stage == \"validation\" else (\"train\", \"val\")\n frames = [\n pd.read_parquet(\n public / \"data\" / f\"{prefix}_likes.parquet\",\n columns=[\"uid\", \"timestamp\", \"item_id\", \"is_organic\"],\n )\n for prefix in prefixes\n ]\n likes = pd.concat(frames, ignore_index=True)\n likes = likes[likes.is_organic == 1].copy()\n mapped = pd.Index(item_ids).get_indexer(\n likes.item_id.to_numpy(dtype=np.int64)\n )\n likes = likes.loc[mapped >= 0, [\"uid\", \"timestamp\"]].copy()\n likes[\"item\"] = mapped[mapped >= 0].astype(np.int32, copy=False)\n likes = likes.sort_values([\"uid\", \"item\", \"timestamp\"], kind=\"stable\")\n nested: dict[int, dict[int, list[int]]] = {}\n for raw_uid, raw_timestamp, raw_item in likes.itertuples(\n index=False, name=None\n ):\n nested.setdefault(int(raw_uid), {}).setdefault(int(raw_item), []).append(\n int(raw_timestamp)\n )\n result = {\n uid: {\n item: np.asarray(timestamps, dtype=np.int64)\n for item, timestamps in items.items()\n }\n for uid, items in nested.items()\n }\n log(\n f\"loaded organic-like outcomes: rows={len(likes):,} \"\n f\"users={len(result):,}\"\n )\n return result\n\n\ndef build_target_weights(\n public: Path,\n stage: str,\n prepared: SequenceData,\n training_counts: np.ndarray,\n config: Config,\n) -> np.ndarray:\n \"\"\"Build nonzero per-token weights without dropping adjacent targets.\"\"\"\n if config.target_weighting == \"uniform\":\n weights = np.ones(len(prepared.sequence_items), dtype=np.float32)\n else:\n weights = np.full(\n len(prepared.sequence_items),\n config.ordinary_target_weight,\n dtype=np.float32,\n )\n organic_strong = (\n (prepared.event_features[:, 0] > 0.5)\n & (prepared.event_features[:, 1] >= 0.8)\n )\n weights[organic_strong] = config.organic_strong_target_weight\n if config.target_weighting == \"discovery-graded\":\n like_outcomes = load_organic_like_outcomes(\n public, stage, prepared.item_ids\n )\n history_end = (\n base.TRAIN_END if stage == \"validation\" else base.VALIDATION_END\n )\n horizon = config.discovery_horizon_days * base.DAY\n grade_three = 0\n for user_row, uid in enumerate(prepared.user_ids):\n start = int(prepared.offsets[user_row])\n end = int(prepared.offsets[user_row + 1])\n if end - start < 2:\n continue\n last_training_target = start + int(training_counts[user_row])\n first_validation_target = last_training_target + 1\n split_timestamp = (\n int(prepared.sequence_timestamps[first_validation_target])\n if first_validation_target < end\n else history_end\n )\n user_likes = like_outcomes.get(int(uid), {})\n seen = {int(prepared.sequence_items[start])}\n for position in range(start + 1, end):\n item = int(prepared.sequence_items[position])\n novel = item not in seen\n seen.add(item)\n if not novel:\n # Repeats remain training examples, but they cannot be\n # eligible discovery outcomes at the corresponding\n # cutoff and therefore keep only the ordinary weight.\n weights[position] = config.ordinary_target_weight\n continue\n timestamps = user_likes.get(item)\n if timestamps is None:\n continue\n timestamp = int(prepared.sequence_timestamps[position])\n upper = min(\n timestamp + horizon,\n split_timestamp\n if position <= last_training_target\n else history_end,\n )\n outcome = int(np.searchsorted(timestamps, timestamp, side=\"left\"))\n if outcome < len(timestamps) and int(timestamps[outcome]) < upper:\n weights[position] = config.organic_like_target_weight\n grade_three += 1\n log(\n f\"grade-3 first-encounter targets={grade_three:,} \"\n f\"horizon_days={config.discovery_horizon_days}\"\n )\n log(\n f\"target weighting={config.target_weighting} \"\n f\"min={weights.min():.3f} mean={weights.mean():.3f} \"\n f\"max={weights.max():.3f}\"\n )\n return weights\n\n\ndef build_dense_block(\n prepared: SequenceData,\n target_weights: np.ndarray,\n users: np.ndarray,\n window_starts: np.ndarray,\n process_counts: np.ndarray,\n score_starts: np.ndarray,\n start: int,\n steps: int,\n config: Config,\n) -> tuple[np.ndarray, ...]:\n \"\"\"Build one causal sequence block and mark the target rows to score.\"\"\"\n batch_size = len(users)\n padding_item = len(prepared.item_ids)\n input_items = np.full((batch_size, steps), padding_item, dtype=np.uint16)\n event_features = np.zeros((batch_size, steps, 3), dtype=np.float16)\n activity_features = np.zeros((batch_size, steps, 4), dtype=np.float16)\n state_ids = np.zeros((batch_size, steps), dtype=np.uint32)\n targets = np.zeros((batch_size, steps), dtype=np.uint16)\n weights = np.ones((batch_size, steps), dtype=np.float32)\n scored = np.zeros((batch_size, steps), dtype=bool)\n histories = np.full(\n (batch_size, steps, config.bptt_steps),\n padding_item,\n dtype=np.uint16,\n )\n processed = np.zeros(batch_size, dtype=np.int32)\n\n lengths = prepared.lengths\n for row, user_row in enumerate(users):\n count = min(steps, max(int(process_counts[row]) - start, 0))\n if not count:\n continue\n offset = int(prepared.offsets[user_row])\n window_start = int(window_starts[row])\n absolute_start = window_start + start\n positions = slice(\n offset + absolute_start,\n offset + absolute_start + count,\n )\n target_positions = slice(\n offset + absolute_start + 1,\n offset + absolute_start + count + 1,\n )\n input_items[row, :count] = prepared.sequence_items[positions]\n event_features[row, :count] = prepared.event_features[positions]\n activity_features[row, :count] = prepared.activity_features[positions]\n state_ids[row, :count] = prepared.state_ids[positions]\n targets[row, :count] = prepared.sequence_items[target_positions]\n weights[row, :count] = target_weights[target_positions]\n score_from = min(count, max(int(score_starts[row]) - start, 0))\n scored[row, score_from:count] = True\n processed[row] = count\n user_sequence = prepared.sequence_items[\n offset : offset + int(lengths[user_row])\n ]\n histories[row, :count] = history_windows(\n user_sequence,\n absolute_start,\n count,\n config.bptt_steps,\n padding_item,\n )\n return (\n input_items,\n event_features,\n activity_features,\n state_ids,\n targets,\n scored,\n histories,\n processed,\n weights,\n )\n\n\ndef score_target_microbatches(\n model: base.SequentialTwoTower,\n prepared: SequenceData,\n config: Config,\n sequence_vectors: torch.Tensor,\n flat_states: np.ndarray,\n flat_activity: np.ndarray,\n positive_items: torch.Tensor,\n negatives: torch.Tensor,\n negative_weights: torch.Tensor,\n sample_weights: np.ndarray,\n device: torch.device,\n backward: bool,\n) -> tuple[dict[str, float], torch.Tensor | None]:\n \"\"\"Compute one logical block's exact BCE in bounded-memory slices.\"\"\"\n target_count = len(positive_items)\n total_weight = float(sample_weights.sum())\n if total_weight <= 0:\n raise RuntimeError(\"target block has no positive sample weight\")\n sequence_gradient = torch.zeros_like(sequence_vectors) if backward else None\n statistics = {\n \"loss_sum\": 0.0,\n \"positive_sum\": 0.0,\n \"negative_sum\": 0.0,\n \"pair_correct\": 0.0,\n \"weight_sum\": 0.0,\n }\n for begin in range(0, target_count, config.loss_microbatch_size):\n end = min(begin + config.loss_microbatch_size, target_count)\n selected_states = flat_states[begin:end]\n if backward:\n sequence_input = (\n sequence_vectors[begin:end].detach().requires_grad_(True)\n )\n else:\n sequence_input = sequence_vectors[begin:end]\n liked_vectors, disliked_vectors, _ = pool_explicit_states(\n model, prepared, selected_states, device\n )\n counts = prepared.state_counts[selected_states]\n numeric = np.zeros((end - begin, 6), dtype=np.float32)\n numeric[:, :4] = flat_activity[begin:end]\n numeric[:, 4] = np.log1p(counts[:, 0]) / 6.0\n numeric[:, 5] = np.log1p(counts[:, 1]) / 4.0\n query_vectors = user_queries(\n model,\n sequence_input,\n liked_vectors,\n disliked_vectors,\n torch.from_numpy(numeric).to(device),\n )\n positive_vectors, negative_vectors = candidate_vectors(\n model,\n positive_items[begin:end],\n negatives[begin:end],\n )\n positive_logits = (query_vectors * positive_vectors).sum(dim=1)\n negative_logits = torch.einsum(\n \"bd,bkd->bk\", query_vectors, negative_vectors\n )\n positive_loss = F.softplus(-positive_logits)\n negative_loss = corrected_negative_loss(\n negative_logits, negative_weights[begin:end], config\n )\n selected_weights = torch.from_numpy(\n sample_weights[begin:end].astype(np.float32, copy=False)\n ).to(device)\n loss_sum = (\n selected_weights * (positive_loss + negative_loss) / 2.0\n ).sum()\n if not bool(torch.isfinite(loss_sum)):\n raise RuntimeError(\"dense next-item training produced non-finite loss\")\n if backward:\n (loss_sum / total_weight).backward()\n if sequence_input.grad is None:\n raise RuntimeError(\"sequence microbatch did not produce a gradient\")\n sequence_gradient[begin:end] = sequence_input.grad\n statistics[\"loss_sum\"] += float(loss_sum.detach().cpu())\n statistics[\"positive_sum\"] += float(positive_logits.detach().sum().cpu())\n statistics[\"negative_sum\"] += float(negative_logits.detach().sum().cpu())\n statistics[\"pair_correct\"] += float(\n (positive_logits.detach()[:, None] > negative_logits.detach()).sum().cpu()\n )\n statistics[\"weight_sum\"] += float(selected_weights.sum().detach().cpu())\n return statistics, sequence_gradient\n\n\ndef release_accelerator_cache(device: torch.device) -> None:\n if device.type == \"mps\":\n torch.mps.empty_cache()\n\n\ndef run_training_epoch(\n model: base.SequentialTwoTower,\n optimizer: torch.optim.Optimizer,\n prepared: SequenceData,\n config: Config,\n target_weights: np.ndarray,\n window_starts: np.ndarray,\n target_counts: np.ndarray,\n item_correction: torch.Tensor,\n known_positive_packed: np.ndarray | None,\n generator: torch.Generator,\n device: torch.device,\n epoch: int,\n) -> dict[str, float]:\n model.train()\n train_users = np.flatnonzero(target_counts > 0)\n sorted_users = train_users[\n np.argsort(target_counts[train_users], kind=\"stable\")\n ]\n groups = [\n sorted_users[start : start + config.user_batch_size]\n for start in range(0, len(sorted_users), config.user_batch_size)\n ]\n ]\n group_order = torch.randperm(len(groups), generator=generator).numpy()\n negative_count = config.uniform_negatives + config.inbatch_negatives\n total_loss = 0.0\n total_targets = 0\n total_weight = 0.0\n total_positive = 0.0\n total_negative = 0.0\n total_pair_correct = 0.0\n batches = 0\n started = time.monotonic()\n\n for group_number in group_order:\n users = groups[int(group_number)]\n group_counts = target_counts[users]\n hidden = torch.zeros(\n 1, len(users), config.embedding_dim, device=device\n )\n max_targets = int(group_counts.max())\n for start in range(0, max_targets, config.bptt_steps):\n steps = min(config.bptt_steps, max_targets - start)\n block = build_dense_block(\n prepared,\n target_weights,\n users,\n window_starts[users],\n group_counts,\n np.zeros(len(users), dtype=np.int64),\n start,\n steps,\n config,\n )\n sequence_outputs, _ = event_sequence(\n model,\n torch.from_numpy(block[0].astype(np.int64)).to(device),\n torch.from_numpy(block[1].astype(np.float32)).to(device),\n hidden,\n )\n hidden = torch.stack(\n [\n sequence_outputs[row, count - 1]\n if count\n else hidden[0, row]\n for row, count in enumerate(block[7])\n ]\n ).unsqueeze(0).detach()\n flat_scored = block[5].reshape(-1)\n if not np.any(flat_scored):\n continue\n mask = torch.from_numpy(flat_scored).to(device)\n sequence_vectors = sequence_outputs.reshape(\n -1, config.embedding_dim\n )[mask]\n flat_states = block[3].reshape(-1)[flat_scored]\n positive_items = torch.from_numpy(\n block[4].reshape(-1)[flat_scored].astype(np.int64)\n ).to(device)\n known_history = torch.from_numpy(\n block[6].reshape(-1, config.bptt_steps)[flat_scored].astype(\n np.int64\n )\n ).to(device)\n liked_items = torch.from_numpy(\n prepared.liked_states[flat_states].astype(np.int64)\n ).to(device)\n flat_user_rows = np.broadcast_to(\n users[:, None], block[5].shape\n ).reshape(-1)[flat_scored]\n negatives, negative_weights = draw_training_negatives(\n positive_items,\n known_history,\n liked_items,\n flat_user_rows,\n known_positive_packed,\n prepared,\n config,\n item_correction,\n generator,\n device,\n )\n optimizer.zero_grad(set_to_none=True)\n statistics, sequence_gradient = score_target_microbatches(\n model,\n prepared,\n config,\n sequence_vectors,\n flat_states,\n block[2].reshape(-1, 4)[flat_scored].astype(np.float32),\n positive_items,\n negatives,\n negative_weights,\n block[8].reshape(-1)[flat_scored],\n device,\n backward=True,\n )\n if sequence_gradient is None:\n raise RuntimeError(\"training block omitted its sequence gradient\")\n sequence_vectors.backward(sequence_gradient)\n nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)\n optimizer.step()\n\n count = len(positive_items)\n total_loss += statistics[\"loss_sum\"]\n total_targets += count\n total_weight += statistics[\"weight_sum\"]\n total_positive += statistics[\"positive_sum\"]\n total_negative += statistics[\"negative_sum\"]\n total_pair_correct += statistics[\"pair_correct\"]\n batches += 1\n if batches % 25 == 0:\n release_accelerator_cache(device)\n if batches % 250 == 0:\n batch_bce = statistics[\"loss_sum\"] / max(\n statistics[\"weight_sum\"], 1.0e-12\n )\n running_bce = total_loss / max(total_weight, 1.0e-12)\n log(\n f\"epoch={epoch}/{config.epochs} batches={batches:,} \"\n f\"targets={total_targets:,} \"\n f\"batch_bce={batch_bce:.6f} \"\n f\"running_bce={running_bce:.6f} \"\n f\"elapsed={time.monotonic() - started:.1f}s\"\n )\n\n return {\n \"epoch\": float(epoch),\n \"loss\": total_loss / max(total_weight, 1.0e-12),\n \"positive_logit\": total_positive / max(total_targets, 1),\n \"negative_logit\": total_negative / max(total_targets * negative_count, 1),\n \"pair_accuracy\": total_pair_correct\n / max(total_targets * negative_count, 1),\n \"targets\": float(total_targets),\n \"target_weight_sum\": total_weight,\n \"batches\": float(batches),\n \"seconds\": time.monotonic() - started,\n }\n\n\[email protected]_grad()\ndef evaluate_validation_bce(\n model: base.SequentialTwoTower,\n prepared: SequenceData,\n config: Config,\n target_weights: np.ndarray,\n window_starts: np.ndarray,\n training_counts: np.ndarray,\n validation_counts: np.ndarray,\n item_correction: torch.Tensor,\n known_positive_packed: np.ndarray | None,\n device: torch.device,\n) -> dict[str, float]:\n \"\"\"Measure deterministic BCE on each user's chronological target suffix.\"\"\"\n model.eval()\n totals = training_counts + validation_counts\n users_with_validation = np.flatnonzero(validation_counts > 0)\n ordered = users_with_validation[\n np.argsort(totals[users_with_validation], kind=\"stable\")\n ]\n groups = [\n ordered[start : start + config.user_batch_size]\n for start in range(0, len(ordered), config.user_batch_size)\n ]\n generator = torch.Generator(device=\"cpu\").manual_seed(config.seed + 20_000)\n total_loss = 0.0\n total_targets = 0\n total_weight = 0.0\n started = time.monotonic()\n\n for group_number, users in enumerate(groups, start=1):\n group_totals = totals[users]\n group_starts = training_counts[users]\n hidden = torch.zeros(\n 1, len(users), config.embedding_dim, device=device\n )\n max_targets = int(group_totals.max())\n for start in range(0, max_targets, config.bptt_steps):\n steps = min(config.bptt_steps, max_targets - start)\n block = build_dense_block(\n prepared,\n target_weights,\n users,\n window_starts[users],\n group_totals,\n group_starts,\n start,\n steps,\n config,\n )\n sequence_outputs, _ = event_sequence(\n model,\n torch.from_numpy(block[0].astype(np.int64)).to(device),\n torch.from_numpy(block[1].astype(np.float32)).to(device),\n hidden,\n )\n hidden = torch.stack(\n [\n sequence_outputs[row, count - 1]\n if count\n else hidden[0, row]\n for row, count in enumerate(block[7])\n ]\n ).unsqueeze(0)\n flat_scored = block[5].reshape(-1)\n if not np.any(flat_scored):\n continue\n mask = torch.from_numpy(flat_scored).to(device)\n sequence_vectors = sequence_outputs.reshape(\n -1, config.embedding_dim\n )[mask]\n flat_states = block[3].reshape(-1)[flat_scored]\n positive_items = torch.from_numpy(\n block[4].reshape(-1)[flat_scored].astype(np.int64)\n ).to(device)\n known_history = torch.from_numpy(\n block[6].reshape(-1, config.bptt_steps)[flat_scored].astype(\n np.int64\n )\n ).to(device)\n liked_items = torch.from_numpy(\n prepared.liked_states[flat_states].astype(np.int64)\n ).to(device)\n flat_user_rows = np.broadcast_to(\n users[:, None], block[5].shape\n ).reshape(-1)[flat_scored]\n negatives, negative_weights = draw_training_negatives(\n positive_items,\n known_history,\n liked_items,\n flat_user_rows,\n known_positive_packed,\n prepared,\n config,\n item_correction,\n generator,\n device,\n )\n statistics, _ = score_target_microbatches(\n model,\n prepared,\n config,\n sequence_vectors,\n flat_states,\n block[2].reshape(-1, 4)[flat_scored].astype(np.float32),\n positive_items,\n negatives,\n negative_weights,\n block[8].reshape(-1)[flat_scored],\n device,\n backward=False,\n )\n total_loss += statistics[\"loss_sum\"]\n total_targets += len(positive_items)\n total_weight += statistics[\"weight_sum\"]\n if group_number % 25 == 0:\n release_accelerator_cache(device)\n return {\n \"validation_loss\": total_loss / max(total_weight, 1.0e-12),\n \"validation_targets\": float(total_targets),\n \"validation_target_weight_sum\": total_weight,\n \"validation_seconds\": time.monotonic() - started,\n }\n\n\ndef snapshot_model(model: nn.Module) -> dict[str, torch.Tensor]:\n return {\n name: value.detach().cpu().clone()\n for name, value in model.state_dict().items()\n }\n\n\ndef save_epoch_checkpoint(\n checkpoint_dir: Path,\n epoch: int,\n model_state: dict[str, torch.Tensor],\n config: Config,\n diagnostics: dict[str, float],\n optimizer: torch.optim.Optimizer,\n generator: torch.Generator,\n) -> Path:\n checkpoint_dir.mkdir(parents=True, exist_ok=True)\n path = checkpoint_dir / f\"epoch-{epoch:02d}.pt\"\n torch.save(\n {\n \"epoch\": epoch,\n \"config\": asdict(config),\n \"diagnostics\": diagnostics,\n \"model_state\": model_state,\n \"optimizer_state\": optimizer.state_dict(),\n \"generator_state\": generator.get_state(),\n },\n path,\n )\n log(f\"saved epoch checkpoint: {path}\")\n return path\n\n\ndef train_model(\n public: Path,\n stage: str,\n prepared: SequenceData,\n config: Config,\n device: torch.device,\n checkpoint_dir: Path | None,\n resume_checkpoint: Path | None,\n) -> tuple[base.SequentialTwoTower, list[dict[str, float]]]:\n partition = chronological_target_partition(\n prepared,\n config.validation_fraction,\n config.max_train_targets_per_user,\n config.max_validation_targets_per_user,\n )\n training_counts = partition.training_counts\n validation_counts = partition.validation_counts\n log(\n f\"chronological BCE window: train={int(training_counts.sum()):,} \"\n f\"validation={int(validation_counts.sum()):,} \"\n f\"selected={partition.selected_count:,} \"\n f\"discarded_old_targets={int(partition.window_starts.sum()):,} \"\n f\"all={prepared.target_count:,}\"\n )\n model = base.SequentialTwoTower(prepared.item_features, model_config(config)).to(\n device\n )\n optimizer = torch.optim.AdamW(\n model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay\n )\n generator = torch.Generator(device=\"cpu\").manual_seed(config.seed + 10_000)\n item_correction = build_item_correction(\n prepared,\n config,\n partition.window_starts,\n training_counts,\n )\n target_weights = build_target_weights(\n public,\n stage,\n prepared,\n partition.history_training_counts,\n config,\n )\n known_positive_packed = (\n build_known_positive_filter(prepared, partition.history_training_counts)\n if config.negative_filter_scope == \"full-training-history\"\n else None\n )\n history: list[dict[str, float]] = []\n best_loss = float(\"inf\")\n best_epoch = 0\n best_state: dict[str, torch.Tensor] | None = None\n epochs_without_improvement = 0\n start_epoch = 0\n\n if resume_checkpoint is not None:\n payload = torch.load(\n resume_checkpoint, map_location=device, weights_only=False\n )\n stored_config = payload.get(\"config\", {})\n architecture_fields = (\n \"bptt_steps\",\n \"embedding_dim\",\n \"hidden_dim\",\n \"user_hidden_1\",\n \"user_hidden_2\",\n \"item_hidden_1\",\n \"item_hidden_2\",\n )\n mismatches = {\n name: (stored_config.get(name), getattr(config, name))\n for name in architecture_fields\n if stored_config.get(name) != getattr(config, name)\n }\n if mismatches:\n raise RuntimeError(\n f\"resume checkpoint architecture differs: {mismatches}\"\n )\n stored_filter_scope = stored_config.get(\n \"negative_filter_scope\", \"last-100\"\n )\n stored_target_weighting = stored_config.get(\n \"target_weighting\", \"uniform\"\n )\n stored_weight_values = (\n stored_config.get(\"ordinary_target_weight\", 1.0),\n stored_config.get(\"organic_strong_target_weight\", 1.0),\n stored_config.get(\"organic_like_target_weight\", 3.0),\n stored_config.get(\"discovery_horizon_days\", 7),\n )\n current_weight_values = (\n config.ordinary_target_weight,\n config.organic_strong_target_weight,\n config.organic_like_target_weight,\n config.discovery_horizon_days,\n )\n stored_window_values = (\n stored_config.get(\"validation_fraction\", 0.10),\n stored_config.get(\"max_train_targets_per_user\", 0),\n stored_config.get(\"max_validation_targets_per_user\", 0),\n )\n current_window_values = (\n config.validation_fraction,\n config.max_train_targets_per_user,\n config.max_validation_targets_per_user,\n )\n objective_changed = (\n stored_filter_scope != config.negative_filter_scope\n or stored_target_weighting != config.target_weighting\n or stored_weight_values != current_weight_values\n or stored_window_values != current_window_values\n )\n model.load_state_dict(payload[\"model_state\"])\n if \"optimizer_state\" in payload and not objective_changed:\n optimizer.load_state_dict(payload[\"optimizer_state\"])\n for group in optimizer.param_groups:\n group[\"lr\"] = config.learning_rate\n group[\"weight_decay\"] = config.weight_decay\n log(f\"restored optimizer state from {resume_checkpoint}\")\n else:\n log(\n f\"using a fresh AdamW optimizer after loading {resume_checkpoint}\"\n )\n if \"generator_state\" in payload and not objective_changed:\n generator.set_state(payload[\"generator_state\"])\n start_epoch = int(payload[\"epoch\"])\n if start_epoch >= config.epochs:\n raise RuntimeError(\n f\"resume epoch {start_epoch} must be below --epochs {config.epochs}\"\n )\n prior = {\n str(name): float(value)\n for name, value in payload.get(\"diagnostics\", {}).items()\n }\n if objective_changed:\n prior.pop(\"validation_loss\", None)\n log(\n \"BCE objective changed: filter \"\n f\"{stored_filter_scope}->{config.negative_filter_scope}, \"\n \"target weighting \"\n f\"{stored_target_weighting}->{config.target_weighting}, \"\n f\"target window {stored_window_values}->{current_window_values}; \"\n \"recomputing the resume checkpoint's validation BCE\"\n )\n if \"validation_loss\" not in prior:\n prior.update(\n evaluate_validation_bce(\n model,\n prepared,\n config,\n target_weights,\n partition.window_starts,\n training_counts,\n validation_counts,\n item_correction,\n known_positive_packed,\n device,\n )\n )\n history.append(prior)\n best_loss = float(prior[\"validation_loss\"])\n best_epoch = start_epoch\n best_state = snapshot_model(model)\n log(\n f\"resumed validation-BCE training: epoch={start_epoch} \"\n f\"loss={best_loss:.6f}\"\n )\n\n for epoch in range(start_epoch + 1, config.epochs + 1):\n diagnostics = run_training_epoch(\n model,\n optimizer,\n prepared,\n config,\n target_weights,\n partition.window_starts,\n training_counts,\n item_correction,\n known_positive_packed,\n generator,\n device,\n epoch,\n )\n validation = evaluate_validation_bce(\n model,\n prepared,\n config,\n target_weights,\n partition.window_starts,\n training_counts,\n validation_counts,\n item_correction,\n known_positive_packed,\n device,\n )\n diagnostics.update(validation)\n history.append(diagnostics)\n current_state = snapshot_model(model)\n if checkpoint_dir is not None:\n save_epoch_checkpoint(\n checkpoint_dir,\n epoch,\n current_state,\n config,\n diagnostics,\n optimizer,\n generator,\n )\n improved = (\n diagnostics[\"validation_loss\"]\n < best_loss - config.early_stopping_min_delta\n )\n if improved:\n best_loss = diagnostics[\"validation_loss\"]\n best_epoch = epoch\n best_state = current_state\n epochs_without_improvement = 0\n if checkpoint_dir is not None:\n torch.save(\n {\n \"epoch\": epoch,\n \"config\": asdict(config),\n \"diagnostics\": diagnostics,\n \"model_state\": current_state,\n \"optimizer_state\": optimizer.state_dict(),\n \"generator_state\": generator.get_state(),\n },\n checkpoint_dir / \"best.pt\",\n )\n else:\n epochs_without_improvement += 1\n log(\n f\"epoch={epoch}/{config.epochs} train_loss={diagnostics['loss']:.5f} \"\n f\"validation_loss={diagnostics['validation_loss']:.5f} \"\n f\"pair_accuracy={diagnostics['pair_accuracy']:.3f} \"\n f\"train_targets={int(diagnostics['targets']):,} \"\n f\"validation_targets={int(diagnostics['validation_targets']):,} \"\n f\"train_seconds={diagnostics['seconds']:.1f} \"\n f\"validation_seconds={diagnostics['validation_seconds']:.1f}\"\n )\n release_accelerator_cache(device)\n if epochs_without_improvement >= config.early_stopping_patience:\n log(\n f\"early stopping after epoch={epoch}; \"\n f\"best_epoch={best_epoch} best_validation_loss={best_loss:.6f}\"\n )\n break\n\n if best_state is None:\n raise RuntimeError(\"training did not produce a valid checkpoint\")\n model.load_state_dict(best_state)\n model.to(device).eval()\n log(\n f\"restored best validation-BCE checkpoint: epoch={best_epoch} \"\n f\"loss={best_loss:.6f}\"\n )\n return model, history\n\n\[email protected]_mode()\ndef inference_queries(\n model: base.SequentialTwoTower,\n prepared: SequenceData,\n config: Config,\n device: torch.device,\n) -> torch.Tensor:\n sequence_rows = prepared.inference_user_rows\n lengths = np.where(\n sequence_rows >= 0, prepared.lengths[np.maximum(sequence_rows, 0)], 0\n sequence_rows >= 0, prepared.lengths[np.maximum(sequence_rows, 0)], 0\n )\n ordered = np.argsort(lengths, kind=\"stable\")\n result = torch.zeros(\n len(sequence_rows), config.embedding_dim, dtype=torch.float32\n )\n for group_start in range(0, len(ordered), config.user_batch_size):\n output_rows = ordered[group_start : group_start + config.user_batch_size]\n rows = sequence_rows[output_rows]\n group_lengths = lengths[output_rows]\n hidden = torch.zeros(\n 1, len(output_rows), config.embedding_dim, device=device\n )\n for start in range(0, int(group_lengths.max(initial=0)), config.bptt_steps):\n steps = min(\n config.bptt_steps,\n int(group_lengths.max(initial=0)) - start,\n )\n items = np.full(\n (len(output_rows), steps),\n len(prepared.item_ids),\n dtype=np.uint16,\n )\n features = np.zeros((len(output_rows), steps, 3), dtype=np.float16)\n counts = np.zeros(len(output_rows), dtype=np.int32)\n for local_row, sequence_row in enumerate(rows):\n if sequence_row < 0:\n continue\n count = min(steps, max(int(group_lengths[local_row]) - start, 0))\n if not count:\n continue\n offset = int(prepared.offsets[sequence_row])\n selected = slice(offset + start, offset + start + count)\n items[local_row, :count] = prepared.sequence_items[selected]\n features[local_row, :count] = prepared.event_features[selected]\n counts[local_row] = count\n outputs, _ = event_sequence(\n model,\n torch.from_numpy(items.astype(np.int64)).to(device),\n torch.from_numpy(features.astype(np.float32)).to(device),\n hidden,\n )\n next_rows: list[torch.Tensor] = []\n for local_row, count in enumerate(counts):\n next_rows.append(\n outputs[local_row, count - 1]\n if count\n else hidden[0, local_row]\n )\n hidden = torch.stack(next_rows).unsqueeze(0)\n\n liked_items = torch.from_numpy(\n prepared.inference_liked_items[output_rows].astype(np.int64)\n ).to(device)\n disliked_items = torch.from_numpy(\n prepared.inference_disliked_items[output_rows].astype(np.int64)\n ).to(device)\n queries = user_queries(\n model,\n hidden.squeeze(0),\n model._mean_context_channel(liked_items),\n model._mean_context_channel(disliked_items),\n torch.from_numpy(\n prepared.inference_user_features[output_rows].astype(np.float32)\n ).to(device),\n )\n result[torch.from_numpy(output_rows.astype(np.int64))] = queries.cpu()\n return result\n\n\[email protected]_mode()\ndef rank_candidates(\n model: base.SequentialTwoTower,\n prepared: SequenceData,\n eligibility: CandidateSet,\n config: Config,\n device: torch.device,\n ranking_batch_size: int,\n) -> tuple[pd.DataFrame, dict[str, float]]:\n model.eval()\n queries = inference_queries(model, prepared, config, device)\n item_parts: list[torch.Tensor] = []\n for start in range(0, len(prepared.item_ids), 2_048):\n rows = torch.arange(\n start,\n min(start + 2_048, len(prepared.item_ids)),\n dtype=torch.long,\n device=device,\n )\n item_parts.append(model.encode_items(rows).cpu())\n items = torch.cat(item_parts)\n rows: list[tuple[int, int, int]] = []\n for start in range(0, len(eligibility.uids), ranking_batch_size):\n end = min(start + ranking_batch_size, len(eligibility.uids))\n scores = queries[start:end] @ items.T\n eligible = np.unpackbits(\n eligibility.packed[start:end],\n axis=1,\n count=len(prepared.item_ids),\n bitorder=\"little\",\n ).astype(bool, copy=False)\n scores.masked_fill_(~torch.from_numpy(eligible), -torch.inf)\n selected = torch.topk(scores, k=base.TOP_N, dim=1).indices.numpy()\n for local_row, columns in enumerate(selected):\n uid = int(eligibility.uids[start + local_row])\n rows.extend(\n (uid, int(prepared.item_ids[column]), rank)\n for rank, column in enumerate(columns, start=1)\n )\n if end == len(eligibility.uids) or end % 1_024 == 0:\n log(f\"ranked users={end:,}/{len(eligibility.uids):,}\")\n query_norm = torch.linalg.vector_norm(queries, dim=1).numpy()\n item_norm = torch.linalg.vector_norm(items, dim=1).numpy()\n diagnostics = {\n \"query_norm_mean\": float(query_norm.mean()),\n \"query_norm_std\": float(query_norm.std()),\n \"item_norm_mean\": float(item_norm.mean()),\n \"item_norm_std\": float(item_norm.std()),\n }\n return pd.DataFrame(rows, columns=[\"uid\", \"item_id\", \"rank\"]), diagnostics\n\n\ndef write_ranking(frame: pd.DataFrame, output: Path) -> None:\n output.parent.mkdir(parents=True, exist_ok=True)\n if output.suffix == \".parquet\":\n frame.to_parquet(output, index=False, compression=\"zstd\")\n elif output.suffix == \".csv\":\n frame.to_csv(output, index=False)\n else:\n raise ValueError(\"--output must end in .parquet or .csv\")\n\n\ndef main() -> int:\n parser = argparse.ArgumentParser(description=__doc__)\n parser.add_argument(\n \"--stage\", choices=(\"validation\", \"submission\"), default=\"validation\"\n )\n parser.add_argument(\"--output\", type=Path, required=True)\n parser.add_argument(\"--public-dir\", type=Path, required=True)\n parser.add_argument(\"--prepared-cache\", type=Path)\n parser.add_argument(\"--device\", choices=(\"cpu\", \"mps\"), default=\"cpu\")\n parser.add_argument(\"--threads\", type=int, default=8)\n parser.add_argument(\"--bptt-steps\", type=int, default=Config.bptt_steps)\n parser.add_argument(\n \"--user-batch-size\", type=int, default=Config.user_batch_size\n )\n parser.add_argument(\"--epochs\", type=int, default=Config.epochs)\n parser.add_argument(\"--learning-rate\", type=float, default=Config.learning_rate)\n parser.add_argument(\n \"--loss-microbatch-size\",\n type=int,\n default=Config.loss_microbatch_size,\n )\n parser.add_argument(\n \"--negative-filter-scope\",\n choices=(\"full-training-history\", \"last-100\"),\n default=Config.negative_filter_scope,\n )\n parser.add_argument(\n \"--target-weighting\",\n choices=(\"discovery-graded\", \"organic-strong\", \"uniform\"),\n default=Config.target_weighting,\n )\n parser.add_argument(\n \"--ordinary-target-weight\",\n type=float,\n default=Config.ordinary_target_weight,\n )\n parser.add_argument(\n \"--organic-strong-target-weight\",\n type=float,\n default=Config.organic_strong_target_weight,\n )\n parser.add_argument(\n \"--organic-like-target-weight\",\n type=float,\n default=Config.organic_like_target_weight,\n )\n parser.add_argument(\n \"--discovery-horizon-days\",\n type=int,\n default=Config.discovery_horizon_days,\n )\n parser.add_argument(\n \"--validation-fraction\",\n type=float,\n default=Config.validation_fraction,\n )\n parser.add_argument(\n \"--max-train-targets-per-user\",\n type=int,\n default=Config.max_train_targets_per_user,\n help=\"retain this many recent pre-validation targets per user; 0 is uncapped\",\n )\n parser.add_argument(\n \"--max-validation-targets-per-user\",\n type=int,\n default=Config.max_validation_targets_per_user,\n help=\"retain this many final validation targets per user; 0 is uncapped\",\n )\n parser.add_argument(\n \"--early-stopping-patience\",\n type=int,\n default=Config.early_stopping_patience,\n )\n parser.add_argument(\n \"--early-stopping-min-delta\",\n type=float,\n default=Config.early_stopping_min_delta,\n )\n parser.add_argument(\"--checkpoint-dir\", type=Path)\n parser.add_argument(\"--resume-checkpoint\", type=Path)\n parser.add_argument(\"--seed\", type=int, default=Config.seed)\n parser.add_argument(\"--ranking-batch-size\", type=int, default=64)\n parser.add_argument(\"--prepare-only\", action=\"store_true\")\n args = parser.parse_args()\n\n if args.threads < 1 or args.bptt_steps < 2 or args.user_batch_size < 1:\n parser.error(\"threads, BPTT steps, and user batch size must be positive\")\n if (\n args.epochs < 1\n or args.learning_rate <= 0\n or args.loss_microbatch_size < 1\n or args.early_stopping_patience < 1\n or args.early_stopping_min_delta < 0\n or args.ordinary_target_weight <= 0\n or args.organic_strong_target_weight <= 0\n or args.organic_like_target_weight <= 0\n or args.discovery_horizon_days < 1\n or args.max_train_targets_per_user < 0\n or args.max_validation_targets_per_user < 0\n ):\n parser.error(\"training counts, learning rate, and patience are invalid\")\n if not 0.0 < args.validation_fraction < 1.0:\n parser.error(\"validation fraction must be strictly between zero and one\")\n public = args.public_dir.resolve()\n if not (public / \"data\" / \"candidate_catalog.parquet\").is_file():\n parser.error(f\"public task data not found under {public}\")\n config = Config(\n bptt_steps=args.bptt_steps,\n user_batch_size=args.user_batch_size,\n epochs=args.epochs,\n learning_rate=args.learning_rate,\n loss_microbatch_size=args.loss_microbatch_size,\n negative_filter_scope=args.negative_filter_scope,\n target_weighting=args.target_weighting,\n ordinary_target_weight=args.ordinary_target_weight,\n organic_strong_target_weight=args.organic_strong_target_weight,\n organic_like_target_weight=args.organic_like_target_weight,\n discovery_horizon_days=args.discovery_horizon_days,\n validation_fraction=args.validation_fraction,\n max_train_targets_per_user=args.max_train_targets_per_user,\n max_validation_targets_per_user=args.max_validation_targets_per_user,\n early_stopping_patience=args.early_stopping_patience,\n early_stopping_min_delta=args.early_stopping_min_delta,\n seed=args.seed,\n )\n torch.set_num_threads(args.threads)\n torch.set_num_interop_threads(1)\n seed_everything(config.seed)\n device = torch.device(args.device)\n if device.type == \"mps\" and not torch.backends.mps.is_available():\n parser.error(\"MPS was requested but is unavailable\")\n\n started = time.monotonic()\n prepared, eligibility = load_or_prepare(\n public,\n args.stage,\n config,\n args.prepared_cache.resolve() if args.prepared_cache else None,\n )\n target_partition = chronological_target_partition(\n prepared,\n config.validation_fraction,\n config.max_train_targets_per_user,\n config.max_validation_targets_per_user,\n )\n if args.prepare_only:\n print(\n json.dumps(\n {\n \"stage\": args.stage,\n \"sequence_users\": len(prepared.user_ids),\n \"sequence_tokens\": len(prepared.sequence_items),\n \"next_item_targets\": prepared.target_count,\n \"selected_training_targets\": int(\n target_partition.training_counts.sum()\n ),\n \"selected_validation_targets\": int(\n target_partition.validation_counts.sum()\n ),\n \"discarded_old_targets\": int(\n target_partition.window_starts.sum()\n ),\n \"target_users\": len(eligibility.uids),\n \"runtime_seconds\": round(time.monotonic() - started, 3),\n \"config\": asdict(config),\n },\n indent=2,\n )\n )\n return 0\n\n model, training_history = train_model(\n public,\n args.stage,\n",
"stderr": "",
"aggregated_output": " )\n\n\ndef event_sequence(\n model: base.SequentialTwoTower,\n items: torch.Tensor,\n features: torch.Tensor,\n hidden: torch.Tensor,\n) -> tuple[torch.Tensor, torch.Tensor]:\n inputs = torch.cat([model.context_embedding(items), features], dim=-1)\n projected = model.event_projection(inputs)\n return model.sequence_encoder(projected, hidden)\n\n\ndef pool_explicit_states(\n model: base.SequentialTwoTower,\n prepared: SequenceData,\n state_ids: np.ndarray,\n device: torch.device,\n) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n unique_states, inverse = np.unique(state_ids, return_inverse=True)\n liked_items = torch.from_numpy(\n prepared.liked_states[unique_states].astype(np.int64)\n ).to(device)\n disliked_items = torch.from_numpy(\n prepared.disliked_states[unique_states].astype(np.int64)\n ).to(device)\n liked_vectors = model._mean_context_channel(liked_items)\n disliked_vectors = model._mean_context_channel(disliked_items)\n inverse_rows = torch.from_numpy(inverse.astype(np.int64)).to(device)\n counts = prepared.state_counts[state_ids]\n return (\n liked_vectors[inverse_rows],\n disliked_vectors[inverse_rows],\n torch.from_numpy(counts.astype(np.float32)).to(device),\n )\n\n\ndef user_queries(\n model: base.SequentialTwoTower,\n sequence_vectors: torch.Tensor,\n liked_vectors: torch.Tensor,\n disliked_vectors: torch.Tensor,\n user_features: torch.Tensor,\n) -> torch.Tensor:\n return model.user_tower(\n torch.cat(\n [sequence_vectors, liked_vectors, disliked_vectors, user_features],\n dim=1,\n )\n )\n\n\ndef candidate_vectors(\n model: base.SequentialTwoTower,\n positives: torch.Tensor,\n negatives: torch.Tensor,\n) -> tuple[torch.Tensor, torch.Tensor]:\n all_items = torch.cat([positives, negatives.reshape(-1)])\n unique_items, inverse = torch.unique(all_items, return_inverse=True)\n encoded = model.encode_items(unique_items)\n positive_vectors = encoded[inverse[: len(positives)]]\n negative_vectors = encoded[inverse[len(positives) :]].reshape(\n len(positives), negatives.shape[1], -1\n )\n return positive_vectors, negative_vectors\n\n\ndef history_windows(\n sequence: np.ndarray,\n start: int,\n count: int,\n width: int,\n padding_item: int,\n) -> np.ndarray:\n base_start = max(0, start - width + 1)\n segment = sequence[base_start : start + count]\n padded = np.concatenate(\n [np.full(width - 1, padding_item, dtype=np.uint16), segment]\n )\n windows = np.lib.stride_tricks.sliding_window_view(padded, width)\n offset = start - base_start\n return windows[offset : offset + count].copy()\n\n\ndef corrected_negative_loss(\n logits: torch.Tensor,\n correction: torch.Tensor,\n config: Config,\n) -> torch.Tensor:\n raw = F.softplus(logits)\n total = config.uniform_negatives + config.inbatch_negatives\n result = torch.zeros(len(logits), device=logits.device)\n if config.uniform_negatives:\n result += config.uniform_negatives * raw[\n :, : config.uniform_negatives\n ].mean(dim=1)\n if config.inbatch_negatives:\n selected = slice(config.uniform_negatives, total)\n weights = correction[:, selected]\n result += config.inbatch_negatives * (\n (raw[:, selected] * weights).sum(dim=1)\n / weights.sum(dim=1).clamp_min(1.0e-6)\n )\n return result / total\n\n\ndef retained_target_window_starts(\n total_targets: np.ndarray,\n validation_counts: np.ndarray,\n training_counts: np.ndarray,\n) -> np.ndarray:\n \"\"\"Return the discarded prefix before the retained chronological suffix.\"\"\"\n training_region_ends = total_targets - validation_counts\n window_starts = training_region_ends - training_counts\n if (\n np.any(window_starts < 0)\n or np.any(training_counts < 0)\n or np.any(validation_counts < 0)\n ):\n raise RuntimeError(\"chronological target partition produced negative counts\")\n if not np.array_equal(\n window_starts + training_counts + validation_counts, total_targets\n ):\n raise RuntimeError(\"chronological target window is not a final contiguous suffix\")\n return window_starts\n\n\ndef chronological_target_partition(\n prepared: SequenceData,\n validation_fraction: float,\n max_train_targets_per_user: int,\n max_validation_targets_per_user: int,\n) -> TargetPartition:\n \"\"\"Select a recent contiguous window with a final out-of-time BCE suffix.\"\"\"\n totals = np.maximum(prepared.lengths - 1, 0).astype(np.int64)\n validation = np.floor(totals * validation_fraction).astype(np.int64)\n validation[(totals >= 2) & (validation == 0)] = 1\n if max_validation_targets_per_user:\n validation = np.minimum(validation, max_validation_targets_per_user)\n\n history_training = totals - validation\n training = history_training.copy()\n if max_train_targets_per_user:\n training = np.minimum(training, max_train_targets_per_user)\n window_starts = retained_target_window_starts(totals, validation, training)\n return TargetPartition(\n window_starts=window_starts,\n training_counts=training,\n validation_counts=validation,\n history_training_counts=history_training,\n )\n\n\ndef load_organic_like_outcomes(\n public: Path, stage: str, item_ids: np.ndarray\n) -> dict[int, dict[int, np.ndarray]]:\n \"\"\"Load mapped organic-like timestamps from public fitting history.\"\"\"\n prefixes = (\"train\",) if stage == \"validation\" else (\"train\", \"val\")\n frames = [\n pd.read_parquet(\n public / \"data\" / f\"{prefix}_likes.parquet\",\n columns=[\"uid\", \"timestamp\", \"item_id\", \"is_organic\"],\n )\n for prefix in prefixes\n ]\n likes = pd.concat(frames, ignore_index=True)\n likes = likes[likes.is_organic == 1].copy()\n mapped = pd.Index(item_ids).get_indexer(\n likes.item_id.to_numpy(dtype=np.int64)\n )\n likes = likes.loc[mapped >= 0, [\"uid\", \"timestamp\"]].copy()\n likes[\"item\"] = mapped[mapped >= 0].astype(np.int32, copy=False)\n likes = likes.sort_values([\"uid\", \"item\", \"timestamp\"], kind=\"stable\")\n nested: dict[int, dict[int, list[int]]] = {}\n for raw_uid, raw_timestamp, raw_item in likes.itertuples(\n index=False, name=None\n ):\n nested.setdefault(int(raw_uid), {}).setdefault(int(raw_item), []).append(\n int(raw_timestamp)\n )\n result = {\n uid: {\n item: np.asarray(timestamps, dtype=np.int64)\n for item, timestamps in items.items()\n }\n for uid, items in nested.items()\n }\n log(\n f\"loaded organic-like outcomes: rows={len(likes):,} \"\n f\"users={len(result):,}\"\n )\n return result\n\n\ndef build_target_weights(\n public: Path,\n stage: str,\n prepared: SequenceData,\n training_counts: np.ndarray,\n config: Config,\n) -> np.ndarray:\n \"\"\"Build nonzero per-token weights without dropping adjacent targets.\"\"\"\n if config.target_weighting == \"uniform\":\n weights = np.ones(len(prepared.sequence_items), dtype=np.float32)\n else:\n weights = np.full(\n len(prepared.sequence_items),\n config.ordinary_target_weight,\n dtype=np.float32,\n )\n organic_strong = (\n (prepared.event_features[:, 0] > 0.5)\n & (prepared.event_features[:, 1] >= 0.8)\n )\n weights[organic_strong] = config.organic_strong_target_weight\n if config.target_weighting == \"discovery-graded\":\n like_outcomes = load_organic_like_outcomes(\n public, stage, prepared.item_ids\n )\n history_end = (\n base.TRAIN_END if stage == \"validation\" else base.VALIDATION_END\n )\n horizon = config.discovery_horizon_days * base.DAY\n grade_three = 0\n for user_row, uid in enumerate(prepared.user_ids):\n start = int(prepared.offsets[user_row])\n end = int(prepared.offsets[user_row + 1])\n if end - start < 2:\n continue\n last_training_target = start + int(training_counts[user_row])\n first_validation_target = last_training_target + 1\n split_timestamp = (\n int(prepared.sequence_timestamps[first_validation_target])\n if first_validation_target < end\n else history_end\n )\n user_likes = like_outcomes.get(int(uid), {})\n seen = {int(prepared.sequence_items[start])}\n for position in range(start + 1, end):\n item = int(prepared.sequence_items[position])\n novel = item not in seen\n seen.add(item)\n if not novel:\n # Repeats remain training examples, but they cannot be\n # eligible discovery outcomes at the corresponding\n # cutoff and therefore keep only the ordinary weight.\n weights[position] = config.ordinary_target_weight\n continue\n timestamps = user_likes.get(item)\n if timestamps is None:\n continue\n timestamp = int(prepared.sequence_timestamps[position])\n upper = min(\n timestamp + horizon,\n split_timestamp\n if position <= last_training_target\n else history_end,\n )\n outcome = int(np.searchsorted(timestamps, timestamp, side=\"left\"))\n if outcome < len(timestamps) and int(timestamps[outcome]) < upper:\n weights[position] = config.organic_like_target_weight\n grade_three += 1\n log(\n f\"grade-3 first-encounter targets={grade_three:,} \"\n f\"horizon_days={config.discovery_horizon_days}\"\n )\n log(\n f\"target weighting={config.target_weighting} \"\n f\"min={weights.min():.3f} mean={weights.mean():.3f} \"\n f\"max={weights.max():.3f}\"\n )\n return weights\n\n\ndef build_dense_block(\n prepared: SequenceData,\n target_weights: np.ndarray,\n users: np.ndarray,\n window_starts: np.ndarray,\n process_counts: np.ndarray,\n score_starts: np.ndarray,\n start: int,\n steps: int,\n config: Config,\n) -> tuple[np.ndarray, ...]:\n \"\"\"Build one causal sequence block and mark the target rows to score.\"\"\"\n batch_size = len(users)\n padding_item = len(prepared.item_ids)\n input_items = np.full((batch_size, steps), padding_item, dtype=np.uint16)\n event_features = np.zeros((batch_size, steps, 3), dtype=np.float16)\n activity_features = np.zeros((batch_size, steps, 4), dtype=np.float16)\n state_ids = np.zeros((batch_size, steps), dtype=np.uint32)\n targets = np.zeros((batch_size, steps), dtype=np.uint16)\n weights = np.ones((batch_size, steps), dtype=np.float32)\n scored = np.zeros((batch_size, steps), dtype=bool)\n histories = np.full(\n (batch_size, steps, config.bptt_steps),\n padding_item,\n dtype=np.uint16,\n )\n processed = np.zeros(batch_size, dtype=np.int32)\n\n lengths = prepared.lengths\n for row, user_row in enumerate(users):\n count = min(steps, max(int(process_counts[row]) - start, 0))\n if not count:\n continue\n offset = int(prepared.offsets[user_row])\n window_start = int(window_starts[row])\n absolute_start = window_start + start\n positions = slice(\n offset + absolute_start,\n offset + absolute_start + count,\n )\n target_positions = slice(\n offset + absolute_start + 1,\n offset + absolute_start + count + 1,\n )\n input_items[row, :count] = prepared.sequence_items[positions]\n event_features[row, :count] = prepared.event_features[positions]\n activity_features[row, :count] = prepared.activity_features[positions]\n state_ids[row, :count] = prepared.state_ids[positions]\n targets[row, :count] = prepared.sequence_items[target_positions]\n weights[row, :count] = target_weights[target_positions]\n score_from = min(count, max(int(score_starts[row]) - start, 0))\n scored[row, score_from:count] = True\n processed[row] = count\n user_sequence = prepared.sequence_items[\n offset : offset + int(lengths[user_row])\n ]\n histories[row, :count] = history_windows(\n user_sequence,\n absolute_start,\n count,\n config.bptt_steps,\n padding_item,\n )\n return (\n input_items,\n event_features,\n activity_features,\n state_ids,\n targets,\n scored,\n histories,\n processed,\n weights,\n )\n\n\ndef score_target_microbatches(\n model: base.SequentialTwoTower,\n prepared: SequenceData,\n config: Config,\n sequence_vectors: torch.Tensor,\n flat_states: np.ndarray,\n flat_activity: np.ndarray,\n positive_items: torch.Tensor,\n negatives: torch.Tensor,\n negative_weights: torch.Tensor,\n sample_weights: np.ndarray,\n device: torch.device,\n backward: bool,\n) -> tuple[dict[str, float], torch.Tensor | None]:\n \"\"\"Compute one logical block's exact BCE in bounded-memory slices.\"\"\"\n target_count = len(positive_items)\n total_weight = float(sample_weights.sum())\n if total_weight <= 0:\n raise RuntimeError(\"target block has no positive sample weight\")\n sequence_gradient = torch.zeros_like(sequence_vectors) if backward else None\n statistics = {\n \"loss_sum\": 0.0,\n \"positive_sum\": 0.0,\n \"negative_sum\": 0.0,\n \"pair_correct\": 0.0,\n \"weight_sum\": 0.0,\n }\n for begin in range(0, target_count, config.loss_microbatch_size):\n end = min(begin + config.loss_microbatch_size, target_count)\n selected_states = flat_states[begin:end]\n if backward:\n sequence_input = (\n sequence_vectors[begin:end].detach().requires_grad_(True)\n )\n else:\n sequence_input = sequence_vectors[begin:end]\n liked_vectors, disliked_vectors, _ = pool_explicit_states(\n model, prepared, selected_states, device\n )\n counts = prepared.state_counts[selected_states]\n numeric = np.zeros((end - begin, 6), dtype=np.float32)\n numeric[:, :4] = flat_activity[begin:end]\n numeric[:, 4] = np.log1p(counts[:, 0]) / 6.0\n numeric[:, 5] = np.log1p(counts[:, 1]) / 4.0\n query_vectors = user_queries(\n model,\n sequence_input,\n liked_vectors,\n disliked_vectors,\n torch.from_numpy(numeric).to(device),\n )\n positive_vectors, negative_vectors = candidate_vectors(\n model,\n positive_items[begin:end],\n negatives[begin:end],\n )\n positive_logits = (query_vectors * positive_vectors).sum(dim=1)\n negative_logits = torch.einsum(\n \"bd,bkd->bk\", query_vectors, negative_vectors\n )\n positive_loss = F.softplus(-positive_logits)\n negative_loss = corrected_negative_loss(\n negative_logits, negative_weights[begin:end], config\n )\n selected_weights = torch.from_numpy(\n sample_weights[begin:end].astype(np.float32, copy=False)\n ).to(device)\n loss_sum = (\n selected_weights * (positive_loss + negative_loss) / 2.0\n ).sum()\n if not bool(torch.isfinite(loss_sum)):\n raise RuntimeError(\"dense next-item training produced non-finite loss\")\n if backward:\n (loss_sum / total_weight).backward()\n if sequence_input.grad is None:\n raise RuntimeError(\"sequence microbatch did not produce a gradient\")\n sequence_gradient[begin:end] = sequence_input.grad\n statistics[\"loss_sum\"] += float(loss_sum.detach().cpu())\n statistics[\"positive_sum\"] += float(positive_logits.detach().sum().cpu())\n statistics[\"negative_sum\"] += float(negative_logits.detach().sum().cpu())\n statistics[\"pair_correct\"] += float(\n (positive_logits.detach()[:, None] > negative_logits.detach()).sum().cpu()\n )\n statistics[\"weight_sum\"] += float(selected_weights.sum().detach().cpu())\n return statistics, sequence_gradient\n\n\ndef release_accelerator_cache(device: torch.device) -> None:\n if device.type == \"mps\":\n torch.mps.empty_cache()\n\n\ndef run_training_epoch(\n model: base.SequentialTwoTower,\n optimizer: torch.optim.Optimizer,\n prepared: SequenceData,\n config: Config,\n target_weights: np.ndarray,\n window_starts: np.ndarray,\n target_counts: np.ndarray,\n item_correction: torch.Tensor,\n known_positive_packed: np.ndarray | None,\n generator: torch.Generator,\n device: torch.device,\n epoch: int,\n) -> dict[str, float]:\n model.train()\n train_users = np.flatnonzero(target_counts > 0)\n sorted_users = train_users[\n np.argsort(target_counts[train_users], kind=\"stable\")\n ]\n groups = [\n sorted_users[start : start + config.user_batch_size]\n for start in range(0, len(sorted_users), config.user_batch_size)\n ]\n ]\n group_order = torch.randperm(len(groups), generator=generator).numpy()\n negative_count = config.uniform_negatives + config.inbatch_negatives\n total_loss = 0.0\n total_targets = 0\n total_weight = 0.0\n total_positive = 0.0\n total_negative = 0.0\n total_pair_correct = 0.0\n batches = 0\n started = time.monotonic()\n\n for group_number in group_order:\n users = groups[int(group_number)]\n group_counts = target_counts[users]\n hidden = torch.zeros(\n 1, len(users), config.embedding_dim, device=device\n )\n max_targets = int(group_counts.max())\n for start in range(0, max_targets, config.bptt_steps):\n steps = min(config.bptt_steps, max_targets - start)\n block = build_dense_block(\n prepared,\n target_weights,\n users,\n window_starts[users],\n group_counts,\n np.zeros(len(users), dtype=np.int64),\n start,\n steps,\n config,\n )\n sequence_outputs, _ = event_sequence(\n model,\n torch.from_numpy(block[0].astype(np.int64)).to(device),\n torch.from_numpy(block[1].astype(np.float32)).to(device),\n hidden,\n )\n hidden = torch.stack(\n [\n sequence_outputs[row, count - 1]\n if count\n else hidden[0, row]\n for row, count in enumerate(block[7])\n ]\n ).unsqueeze(0).detach()\n flat_scored = block[5].reshape(-1)\n if not np.any(flat_scored):\n continue\n mask = torch.from_numpy(flat_scored).to(device)\n sequence_vectors = sequence_outputs.reshape(\n -1, config.embedding_dim\n )[mask]\n flat_states = block[3].reshape(-1)[flat_scored]\n positive_items = torch.from_numpy(\n block[4].reshape(-1)[flat_scored].astype(np.int64)\n ).to(device)\n known_history = torch.from_numpy(\n block[6].reshape(-1, config.bptt_steps)[flat_scored].astype(\n np.int64\n )\n ).to(device)\n liked_items = torch.from_numpy(\n prepared.liked_states[flat_states].astype(np.int64)\n ).to(device)\n flat_user_rows = np.broadcast_to(\n users[:, None], block[5].shape\n ).reshape(-1)[flat_scored]\n negatives, negative_weights = draw_training_negatives(\n positive_items,\n known_history,\n liked_items,\n flat_user_rows,\n known_positive_packed,\n prepared,\n config,\n item_correction,\n generator,\n device,\n )\n optimizer.zero_grad(set_to_none=True)\n statistics, sequence_gradient = score_target_microbatches(\n model,\n prepared,\n config,\n sequence_vectors,\n flat_states,\n block[2].reshape(-1, 4)[flat_scored].astype(np.float32),\n positive_items,\n negatives,\n negative_weights,\n block[8].reshape(-1)[flat_scored],\n device,\n backward=True,\n )\n if sequence_gradient is None:\n raise RuntimeError(\"training block omitted its sequence gradient\")\n sequence_vectors.backward(sequence_gradient)\n nn.utils.clip_grad_norm_(model.parameters(), config.gradient_clip)\n optimizer.step()\n\n count = len(positive_items)\n total_loss += statistics[\"loss_sum\"]\n total_targets += count\n total_weight += statistics[\"weight_sum\"]\n total_positive += statistics[\"positive_sum\"]\n total_negative += statistics[\"negative_sum\"]\n total_pair_correct += statistics[\"pair_correct\"]\n batches += 1\n if batches % 25 == 0:\n release_accelerator_cache(device)\n if batches % 250 == 0:\n batch_bce = statistics[\"loss_sum\"] / max(\n statistics[\"weight_sum\"], 1.0e-12\n )\n running_bce = total_loss / max(total_weight, 1.0e-12)\n log(\n f\"epoch={epoch}/{config.epochs} batches={batches:,} \"\n f\"targets={total_targets:,} \"\n f\"batch_bce={batch_bce:.6f} \"\n f\"running_bce={running_bce:.6f} \"\n f\"elapsed={time.monotonic() - started:.1f}s\"\n )\n\n return {\n \"epoch\": float(epoch),\n \"loss\": total_loss / max(total_weight, 1.0e-12),\n \"positive_logit\": total_positive / max(total_targets, 1),\n \"negative_logit\": total_negative / max(total_targets * negative_count, 1),\n \"pair_accuracy\": total_pair_correct\n / max(total_targets * negative_count, 1),\n \"targets\": float(total_targets),\n \"target_weight_sum\": total_weight,\n \"batches\": float(batches),\n \"seconds\": time.monotonic() - started,\n }\n\n\[email protected]_grad()\ndef evaluate_validation_bce(\n model: base.SequentialTwoTower,\n prepared: SequenceData,\n config: Config,\n target_weights: np.ndarray,\n window_starts: np.ndarray,\n training_counts: np.ndarray,\n validation_counts: np.ndarray,\n item_correction: torch.Tensor,\n known_positive_packed: np.ndarray | None,\n device: torch.device,\n) -> dict[str, float]:\n \"\"\"Measure deterministic BCE on each user's chronological target suffix.\"\"\"\n model.eval()\n totals = training_counts + validation_counts\n users_with_validation = np.flatnonzero(validation_counts > 0)\n ordered = users_with_validation[\n np.argsort(totals[users_with_validation], kind=\"stable\")\n ]\n groups = [\n ordered[start : start + config.user_batch_size]\n for start in range(0, len(ordered), config.user_batch_size)\n ]\n generator = torch.Generator(device=\"cpu\").manual_seed(config.seed + 20_000)\n total_loss = 0.0\n total_targets = 0\n total_weight = 0.0\n started = time.monotonic()\n\n for group_number, users in enumerate(groups, start=1):\n group_totals = totals[users]\n group_starts = training_counts[users]\n hidden = torch.zeros(\n 1, len(users), config.embedding_dim, device=device\n )\n max_targets = int(group_totals.max())\n for start in range(0, max_targets, config.bptt_steps):\n steps = min(config.bptt_steps, max_targets - start)\n block = build_dense_block(\n prepared,\n target_weights,\n users,\n window_starts[users],\n group_totals,\n group_starts,\n start,\n steps,\n config,\n )\n sequence_outputs, _ = event_sequence(\n model,\n torch.from_numpy(block[0].astype(np.int64)).to(device),\n torch.from_numpy(block[1].astype(np.float32)).to(device),\n hidden,\n )\n hidden = torch.stack(\n [\n sequence_outputs[row, count - 1]\n if count\n else hidden[0, row]\n for row, count in enumerate(block[7])\n ]\n ).unsqueeze(0)\n flat_scored = block[5].reshape(-1)\n if not np.any(flat_scored):\n continue\n mask = torch.from_numpy(flat_scored).to(device)\n sequence_vectors = sequence_outputs.reshape(\n -1, config.embedding_dim\n )[mask]\n flat_states = block[3].reshape(-1)[flat_scored]\n positive_items = torch.from_numpy(\n block[4].reshape(-1)[flat_scored].astype(np.int64)\n ).to(device)\n known_history = torch.from_numpy(\n block[6].reshape(-1, config.bptt_steps)[flat_scored].astype(\n np.int64\n )\n ).to(device)\n liked_items = torch.from_numpy(\n prepared.liked_states[flat_states].astype(np.int64)\n ).to(device)\n flat_user_rows = np.broadcast_to(\n users[:, None], block[5].shape\n ).reshape(-1)[flat_scored]\n negatives, negative_weights = draw_training_negatives(\n positive_items,\n known_history,\n liked_items,\n flat_user_rows,\n known_positive_packed,\n prepared,\n config,\n item_correction,\n generator,\n device,\n )\n statistics, _ = score_target_microbatches(\n model,\n prepared,\n config,\n sequence_vectors,\n flat_states,\n block[2].reshape(-1, 4)[flat_scored].astype(np.float32),\n positive_items,\n negatives,\n negative_weights,\n block[8].reshape(-1)[flat_scored],\n device,\n backward=False,\n )\n total_loss += statistics[\"loss_sum\"]\n total_targets += len(positive_items)\n total_weight += statistics[\"weight_sum\"]\n if group_number % 25 == 0:\n release_accelerator_cache(device)\n return {\n \"validation_loss\": total_loss / max(total_weight, 1.0e-12),\n \"validation_targets\": float(total_targets),\n \"validation_target_weight_sum\": total_weight,\n \"validation_seconds\": time.monotonic() - started,\n }\n\n\ndef snapshot_model(model: nn.Module) -> dict[str, torch.Tensor]:\n return {\n name: value.detach().cpu().clone()\n for name, value in model.state_dict().items()\n }\n\n\ndef save_epoch_checkpoint(\n checkpoint_dir: Path,\n epoch: int,\n model_state: dict[str, torch.Tensor],\n config: Config,\n diagnostics: dict[str, float],\n optimizer: torch.optim.Optimizer,\n generator: torch.Generator,\n) -> Path:\n checkpoint_dir.mkdir(parents=True, exist_ok=True)\n path = checkpoint_dir / f\"epoch-{epoch:02d}.pt\"\n torch.save(\n {\n \"epoch\": epoch,\n \"config\": asdict(config),\n \"diagnostics\": diagnostics,\n \"model_state\": model_state,\n \"optimizer_state\": optimizer.state_dict(),\n \"generator_state\": generator.get_state(),\n },\n path,\n )\n log(f\"saved epoch checkpoint: {path}\")\n return path\n\n\ndef train_model(\n public: Path,\n stage: str,\n prepared: SequenceData,\n config: Config,\n device: torch.device,\n checkpoint_dir: Path | None,\n resume_checkpoint: Path | None,\n) -> tuple[base.SequentialTwoTower, list[dict[str, float]]]:\n partition = chronological_target_partition(\n prepared,\n config.validation_fraction,\n config.max_train_targets_per_user,\n config.max_validation_targets_per_user,\n )\n training_counts = partition.training_counts\n validation_counts = partition.validation_counts\n log(\n f\"chronological BCE window: train={int(training_counts.sum()):,} \"\n f\"validation={int(validation_counts.sum()):,} \"\n f\"selected={partition.selected_count:,} \"\n f\"discarded_old_targets={int(partition.window_starts.sum()):,} \"\n f\"all={prepared.target_count:,}\"\n )\n model = base.SequentialTwoTower(prepared.item_features, model_config(config)).to(\n device\n )\n optimizer = torch.optim.AdamW(\n model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay\n )\n generator = torch.Generator(device=\"cpu\").manual_seed(config.seed + 10_000)\n item_correction = build_item_correction(\n prepared,\n config,\n partition.window_starts,\n training_counts,\n )\n target_weights = build_target_weights(\n public,\n stage,\n prepared,\n partition.history_training_counts,\n config,\n )\n known_positive_packed = (\n build_known_positive_filter(prepared, partition.history_training_counts)\n if config.negative_filter_scope == \"full-training-history\"\n else None\n )\n history: list[dict[str, float]] = []\n best_loss = float(\"inf\")\n best_epoch = 0\n best_state: dict[str, torch.Tensor] | None = None\n epochs_without_improvement = 0\n start_epoch = 0\n\n if resume_checkpoint is not None:\n payload = torch.load(\n resume_checkpoint, map_location=device, weights_only=False\n )\n stored_config = payload.get(\"config\", {})\n architecture_fields = (\n \"bptt_steps\",\n \"embedding_dim\",\n \"hidden_dim\",\n \"user_hidden_1\",\n \"user_hidden_2\",\n \"item_hidden_1\",\n \"item_hidden_2\",\n )\n mismatches = {\n name: (stored_config.get(name), getattr(config, name))\n for name in architecture_fields\n if stored_config.get(name) != getattr(config, name)\n }\n if mismatches:\n raise RuntimeError(\n f\"resume checkpoint architecture differs: {mismatches}\"\n )\n stored_filter_scope = stored_config.get(\n \"negative_filter_scope\", \"last-100\"\n )\n stored_target_weighting = stored_config.get(\n \"target_weighting\", \"uniform\"\n )\n stored_weight_values = (\n stored_config.get(\"ordinary_target_weight\", 1.0),\n stored_config.get(\"organic_strong_target_weight\", 1.0),\n stored_config.get(\"organic_like_target_weight\", 3.0),\n stored_config.get(\"discovery_horizon_days\", 7),\n )\n current_weight_values = (\n config.ordinary_target_weight,\n config.organic_strong_target_weight,\n config.organic_like_target_weight,\n config.discovery_horizon_days,\n )\n stored_window_values = (\n stored_config.get(\"validation_fraction\", 0.10),\n stored_config.get(\"max_train_targets_per_user\", 0),\n stored_config.get(\"max_validation_targets_per_user\", 0),\n )\n current_window_values = (\n config.validation_fraction,\n config.max_train_targets_per_user,\n config.max_validation_targets_per_user,\n )\n objective_changed = (\n stored_filter_scope != config.negative_filter_scope\n or stored_target_weighting != config.target_weighting\n or stored_weight_values != current_weight_values\n or stored_window_values != current_window_values\n )\n model.load_state_dict(payload[\"model_state\"])\n if \"optimizer_state\" in payload and not objective_changed:\n optimizer.load_state_dict(payload[\"optimizer_state\"])\n for group in optimizer.param_groups:\n group[\"lr\"] = config.learning_rate\n group[\"weight_decay\"] = config.weight_decay\n log(f\"restored optimizer state from {resume_checkpoint}\")\n else:\n log(\n f\"using a fresh AdamW optimizer after loading {resume_checkpoint}\"\n )\n if \"generator_state\" in payload and not objective_changed:\n generator.set_state(payload[\"generator_state\"])\n start_epoch = int(payload[\"epoch\"])\n if start_epoch >= config.epochs:\n raise RuntimeError(\n f\"resume epoch {start_epoch} must be below --epochs {config.epochs}\"\n )\n prior = {\n str(name): float(value)\n for name, value in payload.get(\"diagnostics\", {}).items()\n }\n if objective_changed:\n prior.pop(\"validation_loss\", None)\n log(\n \"BCE objective changed: filter \"\n f\"{stored_filter_scope}->{config.negative_filter_scope}, \"\n \"target weighting \"\n f\"{stored_target_weighting}->{config.target_weighting}, \"\n f\"target window {stored_window_values}->{current_window_values}; \"\n \"recomputing the resume checkpoint's validation BCE\"\n )\n if \"validation_loss\" not in prior:\n prior.update(\n evaluate_validation_bce(\n model,\n prepared,\n config,\n target_weights,\n partition.window_starts,\n training_counts,\n validation_counts,\n item_correction,\n known_positive_packed,\n device,\n )\n )\n history.append(prior)\n best_loss = float(prior[\"validation_loss\"])\n best_epoch = start_epoch\n best_state = snapshot_model(model)\n log(\n f\"resumed validation-BCE training: epoch={start_epoch} \"\n f\"loss={best_loss:.6f}\"\n )\n\n for epoch in range(start_epoch + 1, config.epochs + 1):\n diagnostics = run_training_epoch(\n model,\n optimizer,\n prepared,\n config,\n target_weights,\n partition.window_starts,\n training_counts,\n item_correction,\n known_positive_packed,\n generator,\n device,\n epoch,\n )\n validation = evaluate_validation_bce(\n model,\n prepared,\n config,\n target_weights,\n partition.window_starts,\n training_counts,\n validation_counts,\n item_correction,\n known_positive_packed,\n device,\n )\n diagnostics.update(validation)\n history.append(diagnostics)\n current_state = snapshot_model(model)\n if checkpoint_dir is not None:\n save_epoch_checkpoint(\n checkpoint_dir,\n epoch,\n current_state,\n config,\n diagnostics,\n optimizer,\n generator,\n )\n improved = (\n diagnostics[\"validation_loss\"]\n < best_loss - config.early_stopping_min_delta\n )\n if improved:\n best_loss = diagnostics[\"validation_loss\"]\n best_epoch = epoch\n best_state = current_state\n epochs_without_improvement = 0\n if checkpoint_dir is not None:\n torch.save(\n {\n \"epoch\": epoch,\n \"config\": asdict(config),\n \"diagnostics\": diagnostics,\n \"model_state\": current_state,\n \"optimizer_state\": optimizer.state_dict(),\n \"generator_state\": generator.get_state(),\n },\n checkpoint_dir / \"best.pt\",\n )\n else:\n epochs_without_improvement += 1\n log(\n f\"epoch={epoch}/{config.epochs} train_loss={diagnostics['loss']:.5f} \"\n f\"validation_loss={diagnostics['validation_loss']:.5f} \"\n f\"pair_accuracy={diagnostics['pair_accuracy']:.3f} \"\n f\"train_targets={int(diagnostics['targets']):,} \"\n f\"validation_targets={int(diagnostics['validation_targets']):,} \"\n f\"train_seconds={diagnostics['seconds']:.1f} \"\n f\"validation_seconds={diagnostics['validation_seconds']:.1f}\"\n )\n release_accelerator_cache(device)\n if epochs_without_improvement >= config.early_stopping_patience:\n log(\n f\"early stopping after epoch={epoch}; \"\n f\"best_epoch={best_epoch} best_validation_loss={best_loss:.6f}\"\n )\n break\n\n if best_state is None:\n raise RuntimeError(\"training did not produce a valid checkpoint\")\n model.load_state_dict(best_state)\n model.to(device).eval()\n log(\n f\"restored best validation-BCE checkpoint: epoch={best_epoch} \"\n f\"loss={best_loss:.6f}\"\n )\n return model, history\n\n\[email protected]_mode()\ndef inference_queries(\n model: base.SequentialTwoTower,\n prepared: SequenceData,\n config: Config,\n device: torch.device,\n) -> torch.Tensor:\n sequence_rows = prepared.inference_user_rows\n lengths = np.where(\n sequence_rows >= 0, prepared.lengths[np.maximum(sequence_rows, 0)], 0\n sequence_rows >= 0, prepared.lengths[np.maximum(sequence_rows, 0)], 0\n )\n ordered = np.argsort(lengths, kind=\"stable\")\n result = torch.zeros(\n len(sequence_rows), config.embedding_dim, dtype=torch.float32\n )\n for group_start in range(0, len(ordered), config.user_batch_size):\n output_rows = ordered[group_start : group_start + config.user_batch_size]\n rows = sequence_rows[output_rows]\n group_lengths = lengths[output_rows]\n hidden = torch.zeros(\n 1, len(output_rows), config.embedding_dim, device=device\n )\n for start in range(0, int(group_lengths.max(initial=0)), config.bptt_steps):\n steps = min(\n config.bptt_steps,\n int(group_lengths.max(initial=0)) - start,\n )\n items = np.full(\n (len(output_rows), steps),\n len(prepared.item_ids),\n dtype=np.uint16,\n )\n features = np.zeros((len(output_rows), steps, 3), dtype=np.float16)\n counts = np.zeros(len(output_rows), dtype=np.int32)\n for local_row, sequence_row in enumerate(rows):\n if sequence_row < 0:\n continue\n count = min(steps, max(int(group_lengths[local_row]) - start, 0))\n if not count:\n continue\n offset = int(prepared.offsets[sequence_row])\n selected = slice(offset + start, offset + start + count)\n items[local_row, :count] = prepared.sequence_items[selected]\n features[local_row, :count] = prepared.event_features[selected]\n counts[local_row] = count\n outputs, _ = event_sequence(\n model,\n torch.from_numpy(items.astype(np.int64)).to(device),\n torch.from_numpy(features.astype(np.float32)).to(device),\n hidden,\n )\n next_rows: list[torch.Tensor] = []\n for local_row, count in enumerate(counts):\n next_rows.append(\n outputs[local_row, count - 1]\n if count\n else hidden[0, local_row]\n )\n hidden = torch.stack(next_rows).unsqueeze(0)\n\n liked_items = torch.from_numpy(\n prepared.inference_liked_items[output_rows].astype(np.int64)\n ).to(device)\n disliked_items = torch.from_numpy(\n prepared.inference_disliked_items[output_rows].astype(np.int64)\n ).to(device)\n queries = user_queries(\n model,\n hidden.squeeze(0),\n model._mean_context_channel(liked_items),\n model._mean_context_channel(disliked_items),\n torch.from_numpy(\n prepared.inference_user_features[output_rows].astype(np.float32)\n ).to(device),\n )\n result[torch.from_numpy(output_rows.astype(np.int64))] = queries.cpu()\n return result\n\n\[email protected]_mode()\ndef rank_candidates(\n model: base.SequentialTwoTower,\n prepared: SequenceData,\n eligibility: CandidateSet,\n config: Config,\n device: torch.device,\n ranking_batch_size: int,\n) -> tuple[pd.DataFrame, dict[str, float]]:\n model.eval()\n queries = inference_queries(model, prepared, config, device)\n item_parts: list[torch.Tensor] = []\n for start in range(0, len(prepared.item_ids), 2_048):\n rows = torch.arange(\n start,\n min(start + 2_048, len(prepared.item_ids)),\n dtype=torch.long,\n device=device,\n )\n item_parts.append(model.encode_items(rows).cpu())\n items = torch.cat(item_parts)\n rows: list[tuple[int, int, int]] = []\n for start in range(0, len(eligibility.uids), ranking_batch_size):\n end = min(start + ranking_batch_size, len(eligibility.uids))\n scores = queries[start:end] @ items.T\n eligible = np.unpackbits(\n eligibility.packed[start:end],\n axis=1,\n count=len(prepared.item_ids),\n bitorder=\"little\",\n ).astype(bool, copy=False)\n scores.masked_fill_(~torch.from_numpy(eligible), -torch.inf)\n selected = torch.topk(scores, k=base.TOP_N, dim=1).indices.numpy()\n for local_row, columns in enumerate(selected):\n uid = int(eligibility.uids[start + local_row])\n rows.extend(\n (uid, int(prepared.item_ids[column]), rank)\n for rank, column in enumerate(columns, start=1)\n )\n if end == len(eligibility.uids) or end % 1_024 == 0:\n log(f\"ranked users={end:,}/{len(eligibility.uids):,}\")\n query_norm = torch.linalg.vector_norm(queries, dim=1).numpy()\n item_norm = torch.linalg.vector_norm(items, dim=1).numpy()\n diagnostics = {\n \"query_norm_mean\": float(query_norm.mean()),\n \"query_norm_std\": float(query_norm.std()),\n \"item_norm_mean\": float(item_norm.mean()),\n \"item_norm_std\": float(item_norm.std()),\n }\n return pd.DataFrame(rows, columns=[\"uid\", \"item_id\", \"rank\"]), diagnostics\n\n\ndef write_ranking(frame: pd.DataFrame, output: Path) -> None:\n output.parent.mkdir(parents=True, exist_ok=True)\n if output.suffix == \".parquet\":\n frame.to_parquet(output, index=False, compression=\"zstd\")\n elif output.suffix == \".csv\":\n frame.to_csv(output, index=False)\n else:\n raise ValueError(\"--output must end in .parquet or .csv\")\n\n\ndef main() -> int:\n parser = argparse.ArgumentParser(description=__doc__)\n parser.add_argument(\n \"--stage\", choices=(\"validation\", \"submission\"), default=\"validation\"\n )\n parser.add_argument(\"--output\", type=Path, required=True)\n parser.add_argument(\"--public-dir\", type=Path, required=True)\n parser.add_argument(\"--prepared-cache\", type=Path)\n parser.add_argument(\"--device\", choices=(\"cpu\", \"mps\"), default=\"cpu\")\n parser.add_argument(\"--threads\", type=int, default=8)\n parser.add_argument(\"--bptt-steps\", type=int, default=Config.bptt_steps)\n parser.add_argument(\n \"--user-batch-size\", type=int, default=Config.user_batch_size\n )\n parser.add_argument(\"--epochs\", type=int, default=Config.epochs)\n parser.add_argument(\"--learning-rate\", type=float, default=Config.learning_rate)\n parser.add_argument(\n \"--loss-microbatch-size\",\n type=int,\n default=Config.loss_microbatch_size,\n )\n parser.add_argument(\n \"--negative-filter-scope\",\n choices=(\"full-training-history\", \"last-100\"),\n default=Config.negative_filter_scope,\n )\n parser.add_argument(\n \"--target-weighting\",\n choices=(\"discovery-graded\", \"organic-strong\", \"uniform\"),\n default=Config.target_weighting,\n )\n parser.add_argument(\n \"--ordinary-target-weight\",\n type=float,\n default=Config.ordinary_target_weight,\n )\n parser.add_argument(\n \"--organic-strong-target-weight\",\n type=float,\n default=Config.organic_strong_target_weight,\n )\n parser.add_argument(\n \"--organic-like-target-weight\",\n type=float,\n default=Config.organic_like_target_weight,\n )\n parser.add_argument(\n \"--discovery-horizon-days\",\n type=int,\n default=Config.discovery_horizon_days,\n )\n parser.add_argument(\n \"--validation-fraction\",\n type=float,\n default=Config.validation_fraction,\n )\n parser.add_argument(\n \"--max-train-targets-per-user\",\n type=int,\n default=Config.max_train_targets_per_user,\n help=\"retain this many recent pre-validation targets per user; 0 is uncapped\",\n )\n parser.add_argument(\n \"--max-validation-targets-per-user\",\n type=int,\n default=Config.max_validation_targets_per_user,\n help=\"retain this many final validation targets per user; 0 is uncapped\",\n )\n parser.add_argument(\n \"--early-stopping-patience\",\n type=int,\n default=Config.early_stopping_patience,\n )\n parser.add_argument(\n \"--early-stopping-min-delta\",\n type=float,\n default=Config.early_stopping_min_delta,\n )\n parser.add_argument(\"--checkpoint-dir\", type=Path)\n parser.add_argument(\"--resume-checkpoint\", type=Path)\n parser.add_argument(\"--seed\", type=int, default=Config.seed)\n parser.add_argument(\"--ranking-batch-size\", type=int, default=64)\n parser.add_argument(\"--prepare-only\", action=\"store_true\")\n args = parser.parse_args()\n\n if args.threads < 1 or args.bptt_steps < 2 or args.user_batch_size < 1:\n parser.error(\"threads, BPTT steps, and user batch size must be positive\")\n if (\n args.epochs < 1\n or args.learning_rate <= 0\n or args.loss_microbatch_size < 1\n or args.early_stopping_patience < 1\n or args.early_stopping_min_delta < 0\n or args.ordinary_target_weight <= 0\n or args.organic_strong_target_weight <= 0\n or args.organic_like_target_weight <= 0\n or args.discovery_horizon_days < 1\n or args.max_train_targets_per_user < 0\n or args.max_validation_targets_per_user < 0\n ):\n parser.error(\"training counts, learning rate, and patience are invalid\")\n if not 0.0 < args.validation_fraction < 1.0:\n parser.error(\"validation fraction must be strictly between zero and one\")\n public = args.public_dir.resolve()\n if not (public / \"data\" / \"candidate_catalog.parquet\").is_file():\n parser.error(f\"public task data not found under {public}\")\n config = Config(\n bptt_steps=args.bptt_steps,\n user_batch_size=args.user_batch_size,\n epochs=args.epochs,\n learning_rate=args.learning_rate,\n loss_microbatch_size=args.loss_microbatch_size,\n negative_filter_scope=args.negative_filter_scope,\n target_weighting=args.target_weighting,\n ordinary_target_weight=args.ordinary_target_weight,\n organic_strong_target_weight=args.organic_strong_target_weight,\n organic_like_target_weight=args.organic_like_target_weight,\n discovery_horizon_days=args.discovery_horizon_days,\n validation_fraction=args.validation_fraction,\n max_train_targets_per_user=args.max_train_targets_per_user,\n max_validation_targets_per_user=args.max_validation_targets_per_user,\n early_stopping_patience=args.early_stopping_patience,\n early_stopping_min_delta=args.early_stopping_min_delta,\n seed=args.seed,\n )\n torch.set_num_threads(args.threads)\n torch.set_num_interop_threads(1)\n seed_everything(config.seed)\n device = torch.device(args.device)\n if device.type == \"mps\" and not torch.backends.mps.is_available():\n parser.error(\"MPS was requested but is unavailable\")\n\n started = time.monotonic()\n prepared, eligibility = load_or_prepare(\n public,\n args.stage,\n config,\n args.prepared_cache.resolve() if args.prepared_cache else None,\n )\n target_partition = chronological_target_partition(\n prepared,\n config.validation_fraction,\n config.max_train_targets_per_user,\n config.max_validation_targets_per_user,\n )\n if args.prepare_only:\n print(\n json.dumps(\n {\n \"stage\": args.stage,\n \"sequence_users\": len(prepared.user_ids),\n \"sequence_tokens\": len(prepared.sequence_items),\n \"next_item_targets\": prepared.target_count,\n \"selected_training_targets\": int(\n target_partition.training_counts.sum()\n ),\n \"selected_validation_targets\": int(\n target_partition.validation_counts.sum()\n ),\n \"discarded_old_targets\": int(\n target_partition.window_starts.sum()\n ),\n \"target_users\": len(eligibility.uids),\n \"runtime_seconds\": round(time.monotonic() - started, 3),\n \"config\": asdict(config),\n },\n indent=2,\n )\n )\n return 0\n\n model, training_history = train_model(\n public,\n args.stage,\n",
"exit_code": 0,
"duration": {
"secs": 0,
"nanos": 11971
},
"formatted_output": "Warning: truncated output (original token count: 12627)\nTotal output lines: 1338\n\n )\n\n\ndef event_sequence(\n model: base.SequentialTwoTower,\n items: torch.Tensor,\n features: torch.Tensor,\n hidden: torch.Tensor,\n) -> tuple[torch.Tensor, torch.Tensor]:\n inputs = torch.cat([model.context_embedding(items), features], dim=-1)\n projected = model.event_projection(inputs)\n return model.sequence_encoder(projected, hidden)\n\n\ndef pool_explicit_states(\n model: base.SequentialTwoTower,\n prepared: SequenceData,\n state_ids: np.ndarray,\n device: torch.device,\n) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n unique_states, inverse = np.unique(state_ids, return_inverse=True)\n liked_items = torch.from_numpy(\n prepared.liked_states[unique_states].astype(np.int64)\n ).to(device)\n disliked_items = torch.from_numpy(\n prepared.disliked_states[unique_states].astype(np.int64)\n ).to(device)\n liked_vectors = model._mean_context_channel(liked_items)\n disliked_vectors = model._mean_context_channel(disliked_items)\n inverse_rows = torch.from_numpy(inverse.astype(np.int64)).to(device)\n counts = prepared.state_counts[state_ids]\n return (\n liked_vectors[inverse_rows],\n disliked_vectors[inverse_rows],\n torch.from_numpy(counts.astype(np.float32)).to(device),\n )\n\n\ndef user_queries(\n model: base.SequentialTwoTower,\n sequence_vectors: torch.Tensor,\n liked_vectors: torch.Tensor,\n disliked_vectors: torch.Tensor,\n user_features: torch.Tensor,\n) -> torch.Tensor:\n return model.user_tower(\n torch.cat(\n [sequence_vectors, liked_vectors, disliked_vectors, user_features],\n dim=1,\n )\n )\n\n\ndef candidate_vectors(\n model: base.SequentialTwoTower,\n positives: torch.Tensor,\n negatives: torch.Tensor,\n) -> tuple[torch.Tensor, torch.Tensor]:\n all_items = torch.cat([positives, negatives.reshape(-1)])\n unique_items, inverse = torch.unique(all_items, return_inverse=True)\n encoded = model.encode_items(unique_items)\n positive_vectors = encoded[inverse[: len(positives)]]\n negative_vectors = encoded[inverse[len(positives) :]].reshape(\n len(positives), negatives.shape[1], -1\n )\n return positive_vectors, negative_vectors\n\n\ndef history_windows(\n sequence: np.ndarray,\n start: int,\n count: int,\n width: int,\n padding_item: int,\n) -> np.ndarray:\n base_start = max(0, start - width + 1)\n segment = sequence[base_start : start + count]\n padded = np.concatenate(\n [np.full(width - 1, padding_item, dtype=np.uint16), segment]\n )\n windows = np.lib.stride_tricks.sliding_window_view(padded, width)\n offset = start - base_start\n return windows[offset : offset + count].copy()\n\n\ndef corrected_negative_loss(\n logits: torch.Tensor,\n correction: torch.Tensor,\n config: Config,\n) -> torch.Tensor:\n raw = F.softplus(logits)\n total = config.uniform_negatives + config.inbatch_negatives\n result = torch.zeros(len(logits), device=logits.device)\n if config.uniform_negatives:\n result += config.uniform_negatives * raw[\n :, : config.uniform_negatives\n ].mean(dim=1)\n if config.inbatch_negatives:\n selected = slice(config.uniform_negatives, total)\n weights = correction[:, selected]\n result += config.inbatch_negatives * (\n (raw[:, selected] * weights).sum(dim=1)\n / weights.sum(dim=1).clamp_min(1.0e-6)\n )\n return result / total\n\n\ndef retained_target_window_starts(\n total_targets: np.ndarray,\n validation_counts: np.ndarray,\n training_counts: np.ndarray,\n) -> np.ndarray:\n \"\"\"Return the discarded prefix before the retained chronological suffix.\"\"\"\n training_region_ends = total_targets - validation_counts\n window_starts = training_region_ends - training_counts\n if (\n np.any(window_starts < 0)\n or np.any(training_counts < 0)\n or np.any(validation_counts < 0)\n ):\n raise RuntimeError(\"chronological target partition produced negative counts\")\n if not np.array_equal(\n window_starts + training_counts + validation_counts, total_targets\n ):\n raise RuntimeError(\"chronological target window is not a final contiguous suffix\")\n return window_starts\n\n\ndef chronological_target_partition(\n prepared: SequenceData,\n validation_fraction: float,\n max_train_targets_per_user: int,\n max_validation_targets_per_user: int,\n) -> TargetPartition:\n \"\"\"Select a recent contiguous window with a final out-of-time BCE suffix.\"\"\"\n totals = np.maximum(prepared.lengths - 1, 0).astype(np.int64)\n validation = np.floor(totals * validation_fraction).astype(np.int64)\n validation[(totals >= 2) & (validation == 0)] = 1\n if max_validation_targets_per_user:\n validation = np.minimum(validation, max_validation_targets_per_user)\n\n history_training = totals - validation\n training = history_training.copy()\n if max_train_targets_per_user:\n training = np.minimum(training, max_train_targets_per_user)\n window_starts = retained_target_window_starts(totals, validation, training)\n return TargetPartition(\n window_starts=window_starts,\n training_counts=training,\n validation_counts=validation,\n history_training_counts=history_training,\n )\n\n\ndef load_organic_like_outcomes(\n public: Path, stage: str, item_ids: np.ndarray\n) -> dict[int, dict[int, np.ndarray]]:\n \"\"\"Load mapped organic-like timestamps from public fitting history.\"\"\"\n prefixes = (\"train\",) if stage == \"validation\" else (\"train\", \"val\")\n frames = [\n pd.read_parquet(\n public / \"data\" / f\"{prefix}_likes.parquet\",\n columns=[\"uid\", \"timestamp\", \"item_id\", \"is_organic\"],\n )\n for prefix in prefixes\n ]\n likes = pd.concat(frames, ignore_index=True)\n likes = likes[likes.is_organic == 1].copy()\n mapped = pd.Index(item_ids).get_indexer(\n likes.item_id.to_numpy(dtype=np.int64)\n )\n likes = likes.loc[mapped >= 0, [\"uid\", \"timestamp\"]].copy()\n likes[\"item\"] = mapped[mapped >= 0].astype(np.int32, copy=False)\n likes = likes.sort_values([\"uid\", \"item\", \"timestamp\"], kind=\"stable\")\n nested: dict[int, dict[int, list[int]]] = {}\n for raw_uid, raw_timestamp, raw_item in likes.itertuples(\n index=False, name=None\n ):\n nested.setdefault(int(raw_uid), {}).setdefault(int(raw_item), []).append(\n int(raw_timestamp)\n )\n result = {\n uid: {\n item: np.asarray(timestamps, dtype=np.int64)\n for item, timestamps in items.items()\n }\n for uid, items in nested.items()\n }\n log(\n f\"loaded organic-like outcomes: rows={len(likes):,} \"\n f\"users={len(result):,}\"\n )\n return result\n\n\ndef build_target_weights(\n public: Path,\n stage: str,\n prepared: SequenceData,\n training_counts: np.ndarray,\n config: Config,\n) -> np.ndarray:\n \"\"\"Build nonzero per-token weights without dropping adjacent targets.\"\"\"\n if config.target_weighting == \"uniform\":\n weights = np.ones(len(prepared.sequence_items), dtype=np.float32)\n else:\n weights = np.full(\n len(prepared.sequence_items),\n config.ordinary_target_weight,\n dtype=np.float32,\n )\n organic_strong = (\n (prepared.event_features[:, 0] > 0.5)\n & (prepared.event_features[:, 1] >= 0.8)\n )\n weights[organic_strong] = config.organic_strong_target_weight\n if config.target_weighting == \"discovery-graded\":\n like_outcomes = load_organic_like_outcomes(\n public, stage, prepared.item_ids\n )\n history_end = (\n base.TRAIN_END if stage == \"validation\" else base.VALIDATION_END\n )\n horizon = config.discovery_horizon_days * base.DAY\n grade_three = 0\n for user_row, uid in enumerate(prepared.user_ids):\n start = int(prepared.offsets[user_row])\n end = int(prepared.offsets[user_row + 1])\n if end - start < 2:\n continue\n last_training_target = start + int(training_counts[user_row])\n first_validation_target = last_training_target + 1\n split_timestamp = (\n int(prepared.sequence_timestamps[first_validation_target])\n if first_validation_target < end\n else history_end\n )\n user_likes = like_outcomes.get(int(uid), {})\n seen = {int(prepared.sequence_items[start])}\n for position in range(start + 1, end):\n item = int(prepared.sequence_items[position])\n novel = item not in seen\n seen.add(item)\n if not novel:\n # Repeats remain training examples, but they cannot be\n # eligible discovery outcomes at the corresponding\n # cutoff and therefore keep only the ordinary weight.\n weights[position] = config.ordinary_target_weight\n continue\n timestamps = user_likes.get(item)\n if timestamps is None:\n continue\n timestamp = int(prepared.sequence_timestamps[position])\n upper = min(\n timestamp + horizon,\n split_timestamp\n if position <= last_training_target\n else history_end,\n )\n outcome = int(np.searchsorted(timestamps, timestamp, side=\"left\"))\n if outcome < len(timestamps) and int(timestamps[outcome]) < upper:\n weights[position] = config.organic_like_target_weight\n grade_three += 1\n log(\n f\"grade-3 first-encounter targets={grade_three:,} \"\n f\"horizon_days={config.discovery_horizon_days}\"\n )\n log(\n f\"target weighting={config.target_weighting} \"\n f\"min={weights.min():.3f} mean={weights.mean():.3f} \"\n f\"max={weights.max():.3f}\"\n )\n return weights\n\n\ndef build_dense_block(\n prepared: SequenceData,\n target_weights: np.ndarray,\n users: np.ndarray,\n window_starts: np.ndarray,\n process_counts: np.ndarray,\n score_starts: np.ndarray,\n start: int,\n steps: int,\n config: Config,\n) -> tuple[np.ndarray, ...]:\n \"\"\"Build one causal sequence block and mark the target rows to score.\"\"\"\n batch_size = len(users)\n padding_item = len(prepared.item_ids)\n input_items = np.full((batch_size, steps), padding_item, dtype=np.uint16)\n event_features = np.zeros((batch_size, steps, 3), dtype=np.float16)\n activity_features = np.zeros((batch_size, steps, 4), dtype=np.float16)\n state_ids = np.zeros((batch_size, steps), dtype=np.uint32)\n targets = np.zeros((batch_size, steps), dtype=np.uint16)\n weights = np.ones((batch_size, steps), dtype=np.float32)\n scored = np.zeros((batch_size, steps), dtype=bool)\n histories = np.full(\n (batch_size, steps, config.bptt_steps),\n padding_item,\n dtype=np.uint16,\n )\n processed = np.zeros(batch_size, dtype=np.int32)\n\n lengths = prepared.lengths\n for row, user_row in enumerate(users):\n count = min(steps, max(int(process_counts[row]) - start, 0))\n if not count:\n continue\n offset = int(prepared.offsets[user_row])\n window_start = int(window_starts[row])\n absolute_start = window_start + start\n positions = slice(\n offset + absolute_start,\n offset + absolute_start + count,\n )\n target_positions = slice(\n offset + absolute_start + 1,\n offset + absolute_start + count + 1,\n )\n input_items[row, :count] = prepared.sequence_items[positions]\n event_features[row, :count] = prepared.event_features[positions]\n activity_features[row, :count] = prepared.activity_features[positions]\n state_ids[row, :count] = prepared.state_ids[positions]\n targets[row, :count] = prepared.sequence_items[target_positions]\n weights[row, :count] = target_weights[target_positions]\n score_from = min(count, max(int(score_starts[row]) - start, 0))\n scored[row, score_from:count] = True\n processed[row] = count\n user_sequence = prepared.sequence_items[\n offset : offset + int(lengths[user_row])\n ]\n histories[row, :count] = history_windows(\n user_sequence,\n absolute_start,\n count,\n config.bptt_steps,\n padding_item,\n )\n return (\n input_items,\n event_features,\n activity_features,\n state_ids,\n targets,\n scored,\n histories,\n processed,\n weights,\n )\n\n\ndef score_target_microbatches(\n model: base.SequentialTwoTower,\n prepared: SequenceData,\n config: Config,\n sequence_vectors: torch.Tensor,\n flat_states: np.ndarray,\n flat_activity: np.ndarray,\n positive_items: torch.Tensor,\n negatives: torch.Tensor,\n negative_weights: torch.Tensor,\n sample_weights: np.ndarray,\n device: torch.device,\n backward: bool,\n) -> tuple[dict[str, float], torch.Tensor | None]:\n \"\"\"Compute one logical block's exact BCE in bounded-memory slices.\"\"\"\n target_count = len(positive_items)\n total_weight = float(sample_weights.sum())\n if total_weight <= 0:\n raise RuntimeError(\"target block has no positive sample weight\")\n sequence_gradient = torch.zeros_like(sequence_vectors) if backward else None\n statistics = {\n \"loss_sum\": 0.0,\n \"positive_sum\": 0.0,\n \"negative_sum\": 0.0,\n \"pair_correct\": 0.0,\n \"weight_sum\": 0.0,\n }\n for begin in range(0, target_count, config.loss_microbatch_size):\n end = min(begin + config.loss_microbatch_size, target_count)\n selected_states = flat_states[begin:end]\n if backward:\n sequence_input = (\n sequence_vectors[begin:end].detach().requires_grad_(True)\n )\n else:\n sequence_input = sequence_vectors[begin:end]\n liked_vectors, disliked_vectors, _ = pool_explicit_states(\n model, prepared, selected_states, device\n )\n counts = prepared.state_counts[selected_states]\n numeric = np.zeros((end - begin, 6), dtype=np.float32)\n numeric[:, :4] = flat_activity[begin:end]\n numeric[:, 4] = np.log1p(counts[:, 0]) / 6.0\n numeric[:, 5] = np.log1p(counts[:, 1]) / 4.0\n query_vectors = user_queries(\n model,\n sequence_input,\n liked_vectors,\n disliked_vectors,\n torch.from_numpy(numeric).to(device),\n )\n positive_vectors, negative_vectors = candidate_vectors(\n model,\n positive_items[begin:end],\n negatives[begin:end],\n )\n positive_logits = (query_vectors * positive_vectors).sum(dim=1)\n negative_logits = torch.einsum(\n \"bd,bkd->bk\", query_vectors, negative_vectors\n )\n positive_loss = F.softplus(-positive_logits)\n negative_loss = corrected_negative_loss(\n negative_logits, negative_weights[begin:end], config\n )\n selected_weights = torch.from_numpy(\n sample_weights[begin:end].astype(np.float32, copy=False)\n ).to(device)\n loss_sum = (\n selected_weights * (positive_loss + negative_loss) / 2.0\n ).sum()\n if not bool(torch.isfinite(loss_sum)):\n raise RuntimeError(\"dense next-item training produced non-finite loss\")\n if backward:\n (loss_sum / total_weight).backward()\n if sequence_input.grad is None:\n raise RuntimeError(\"sequence microbatch did not produce a gradient\")\n sequence_gradient[begin:end] = sequence_input.grad\n statistics[\"loss_sum\"] += float(loss_sum.detach().cpu())\n statistics[\"positive_sum\"] += float(positive_logits.detach().sum().cpu())\n statistics[\"negative_sum\"] += float(negative_logits.detach().sum().cpu())\n statistics[\"pair_correct\"] += float(\n (positive_logits.detach()[:, None] > negative_logits.detach()).sum().cpu()\n )\n statistics[\"weight_sum\"] += float(selected_weights.sum().detach().cpu())\n return statistics, sequence_gradient\n\n\ndef release_accelerator_cache(device: torch.device) -> None:\n if device.type == \"mps\":\n torch.mps.empty_cache()\n\n\ndef run_training_epoch(\n model: base.SequentialTwoTower,\n optimizer: torch.optim.Optimizer,\n prepared: SequenceData,\n config: Config,\n target_weights: np.ndarray,\n window_starts: np.ndarray,\n target_counts: np.ndarray,\n item_correction: torch.Tensor,\n known_positive_packed: np.ndarray | None,\n generator: torch.Generator,\n device: torch.device,\n epoch: int,\n) -> dict[str, float]:\n model.train()\n train_users = np.flatnonzero(target_counts > 0)\n sorted_users = train_users[\n np.argsort(target_counts[train_users], kind=\"stable\")\n ]\n groups = [\n sorted_users[start : start + config.user_batch_size]\n for start in range(0, len(sorted_users), config.user_batch_size)\n ]\n ]\n group_order = torch.randperm(len(groups), generator=generator).numpy()\n negative_count = config.uniform_negatives + config.inbatch_negatives\n total_loss = 0.0\n total_targets = 0\n total_weight = 0.0\n total_positive = 0.0\n total_negative = 0.0\n total_pair_correct = 0.0\n batches = 0\n started = time.monotonic()\n\n for group_number in group_order:\n users = groups[int(group_number)]\n group_counts = target_counts[users]\n hidden = torch.zeros(\n 1, len(users), config.embedding_dim, device=device\n )\n max_targets = int(group_counts.max())\n for start in range(0, max_targets, config.bptt_steps):\n steps = min(config.bptt_steps, max_targets - start)\n block = build_dense_block(\n prepared,\n target_weights,\n users,\n window_starts[users],\n group_counts,\n np.zeros(len(users), dtype=np.int64),\n start,\n steps,\n config,\n )\n sequence_outputs, _ = event_sequence(\n model,\n torch.from_numpy(block[0].astype(np.int64)).to(device),\n torch.from_numpy(block[1].astype(np.float32)).to(device),\n hidden,\n )\n hidden = torch.stack(\n [\n sequence_outputs[row, count - 1]\n if count\n else hidden[0, row]\n for row, count in enumerate(block[7])\n ]\n ).unsqueeze(0).detach()\n flat_scored = block[5].reshape(-1)\n if not np.any(flat_scored):\n continue\n mask = torch.from_numpy(flat_scored).to(device)\n sequence_vectors = sequence_outputs.reshape(\n -1, config.embedding_dim\n )[mask]\n flat_states = block[3].reshape(-1)[flat_scored]\n positive_items = torch.from_numpy(\n block[4].reshape(-1)[flat_scored].astype(np.int64)\n ).to(device)\n known_history = torch.from_numpy(\n block[6].reshape(-1, config.bptt_steps)[flat_scored].astype(\n np.int64\n )\n ).to(device)\n liked_items = torch.from_numpy(\n prepared.liked_states[flat_states].astype(np.int64)\n ).to(device)\n flat_user_rows = np.broadcast_to(\n users[:, None], block[5…2627 tokens truncated…dden_1\",\n \"user_hidden_2\",\n \"item_hidden_1\",\n \"item_hidden_2\",\n )\n mismatches = {\n name: (stored_config.get(name), getattr(config, name))\n for name in architecture_fields\n if stored_config.get(name) != getattr(config, name)\n }\n if mismatches:\n raise RuntimeError(\n f\"resume checkpoint architecture differs: {mismatches}\"\n )\n stored_filter_scope = stored_config.get(\n \"negative_filter_scope\", \"last-100\"\n )\n stored_target_weighting = stored_config.get(\n \"target_weighting\", \"uniform\"\n )\n stored_weight_values = (\n stored_config.get(\"ordinary_target_weight\", 1.0),\n stored_config.get(\"organic_strong_target_weight\", 1.0),\n stored_config.get(\"organic_like_target_weight\", 3.0),\n stored_config.get(\"discovery_horizon_days\", 7),\n )\n current_weight_values = (\n config.ordinary_target_weight,\n config.organic_strong_target_weight,\n config.organic_like_target_weight,\n config.discovery_horizon_days,\n )\n stored_window_values = (\n stored_config.get(\"validation_fraction\", 0.10),\n stored_config.get(\"max_train_targets_per_user\", 0),\n stored_config.get(\"max_validation_targets_per_user\", 0),\n )\n current_window_values = (\n config.validation_fraction,\n config.max_train_targets_per_user,\n config.max_validation_targets_per_user,\n )\n objective_changed = (\n stored_filter_scope != config.negative_filter_scope\n or stored_target_weighting != config.target_weighting\n or stored_weight_values != current_weight_values\n or stored_window_values != current_window_values\n )\n model.load_state_dict(payload[\"model_state\"])\n if \"optimizer_state\" in payload and not objective_changed:\n optimizer.load_state_dict(payload[\"optimizer_state\"])\n for group in optimizer.param_groups:\n group[\"lr\"] = config.learning_rate\n group[\"weight_decay\"] = config.weight_decay\n log(f\"restored optimizer state from {resume_checkpoint}\")\n else:\n log(\n f\"using a fresh AdamW optimizer after loading {resume_checkpoint}\"\n )\n if \"generator_state\" in payload and not objective_changed:\n generator.set_state(payload[\"generator_state\"])\n start_epoch = int(payload[\"epoch\"])\n if start_epoch >= config.epochs:\n raise RuntimeError(\n f\"resume epoch {start_epoch} must be below --epochs {config.epochs}\"\n )\n prior = {\n str(name): float(value)\n for name, value in payload.get(\"diagnostics\", {}).items()\n }\n if objective_changed:\n prior.pop(\"validation_loss\", None)\n log(\n \"BCE objective changed: filter \"\n f\"{stored_filter_scope}->{config.negative_filter_scope}, \"\n \"target weighting \"\n f\"{stored_target_weighting}->{config.target_weighting}, \"\n f\"target window {stored_window_values}->{current_window_values}; \"\n \"recomputing the resume checkpoint's validation BCE\"\n )\n if \"validation_loss\" not in prior:\n prior.update(\n evaluate_validation_bce(\n model,\n prepared,\n config,\n target_weights,\n partition.window_starts,\n training_counts,\n validation_counts,\n item_correction,\n known_positive_packed,\n device,\n )\n )\n history.append(prior)\n best_loss = float(prior[\"validation_loss\"])\n best_epoch = start_epoch\n best_state = snapshot_model(model)\n log(\n f\"resumed validation-BCE training: epoch={start_epoch} \"\n f\"loss={best_loss:.6f}\"\n )\n\n for epoch in range(start_epoch + 1, config.epochs + 1):\n diagnostics = run_training_epoch(\n model,\n optimizer,\n prepared,\n config,\n target_weights,\n partition.window_starts,\n training_counts,\n item_correction,\n known_positive_packed,\n generator,\n device,\n epoch,\n )\n validation = evaluate_validation_bce(\n model,\n prepared,\n config,\n target_weights,\n partition.window_starts,\n training_counts,\n validation_counts,\n item_correction,\n known_positive_packed,\n device,\n )\n diagnostics.update(validation)\n history.append(diagnostics)\n current_state = snapshot_model(model)\n if checkpoint_dir is not None:\n save_epoch_checkpoint(\n checkpoint_dir,\n epoch,\n current_state,\n config,\n diagnostics,\n optimizer,\n generator,\n )\n improved = (\n diagnostics[\"validation_loss\"]\n < best_loss - config.early_stopping_min_delta\n )\n if improved:\n best_loss = diagnostics[\"validation_loss\"]\n best_epoch = epoch\n best_state = current_state\n epochs_without_improvement = 0\n if checkpoint_dir is not None:\n torch.save(\n {\n \"epoch\": epoch,\n \"config\": asdict(config),\n \"diagnostics\": diagnostics,\n \"model_state\": current_state,\n \"optimizer_state\": optimizer.state_dict(),\n \"generator_state\": generator.get_state(),\n },\n checkpoint_dir / \"best.pt\",\n )\n else:\n epochs_without_improvement += 1\n log(\n f\"epoch={epoch}/{config.epochs} train_loss={diagnostics['loss']:.5f} \"\n f\"validation_loss={diagnostics['validation_loss']:.5f} \"\n f\"pair_accuracy={diagnostics['pair_accuracy']:.3f} \"\n f\"train_targets={int(diagnostics['targets']):,} \"\n f\"validation_targets={int(diagnostics['validation_targets']):,} \"\n f\"train_seconds={diagnostics['seconds']:.1f} \"\n f\"validation_seconds={diagnostics['validation_seconds']:.1f}\"\n )\n release_accelerator_cache(device)\n if epochs_without_improvement >= config.early_stopping_patience:\n log(\n f\"early stopping after epoch={epoch}; \"\n f\"best_epoch={best_epoch} best_validation_loss={best_loss:.6f}\"\n )\n break\n\n if best_state is None:\n raise RuntimeError(\"training did not produce a valid checkpoint\")\n model.load_state_dict(best_state)\n model.to(device).eval()\n log(\n f\"restored best validation-BCE checkpoint: epoch={best_epoch} \"\n f\"loss={best_loss:.6f}\"\n )\n return model, history\n\n\[email protected]_mode()\ndef inference_queries(\n model: base.SequentialTwoTower,\n prepared: SequenceData,\n config: Config,\n device: torch.device,\n) -> torch.Tensor:\n sequence_rows = prepared.inference_user_rows\n lengths = np.where(\n sequence_rows >= 0, prepared.lengths[np.maximum(sequence_rows, 0)], 0\n sequence_rows >= 0, prepared.lengths[np.maximum(sequence_rows, 0)], 0\n )\n ordered = np.argsort(lengths, kind=\"stable\")\n result = torch.zeros(\n len(sequence_rows), config.embedding_dim, dtype=torch.float32\n )\n for group_start in range(0, len(ordered), config.user_batch_size):\n output_rows = ordered[group_start : group_start + config.user_batch_size]\n rows = sequence_rows[output_rows]\n group_lengths = lengths[output_rows]\n hidden = torch.zeros(\n 1, len(output_rows), config.embedding_dim, device=device\n )\n for start in range(0, int(group_lengths.max(initial=0)), config.bptt_steps):\n steps = min(\n config.bptt_steps,\n int(group_lengths.max(initial=0)) - start,\n )\n items = np.full(\n (len(output_rows), steps),\n len(prepared.item_ids),\n dtype=np.uint16,\n )\n features = np.zeros((len(output_rows), steps, 3), dtype=np.float16)\n counts = np.zeros(len(output_rows), dtype=np.int32)\n for local_row, sequence_row in enumerate(rows):\n if sequence_row < 0:\n continue\n count = min(steps, max(int(group_lengths[local_row]) - start, 0))\n if not count:\n continue\n offset = int(prepared.offsets[sequence_row])\n selected = slice(offset + start, offset + start + count)\n items[local_row, :count] = prepared.sequence_items[selected]\n features[local_row, :count] = prepared.event_features[selected]\n counts[local_row] = count\n outputs, _ = event_sequence(\n model,\n torch.from_numpy(items.astype(np.int64)).to(device),\n torch.from_numpy(features.astype(np.float32)).to(device),\n hidden,\n )\n next_rows: list[torch.Tensor] = []\n for local_row, count in enumerate(counts):\n next_rows.append(\n outputs[local_row, count - 1]\n if count\n else hidden[0, local_row]\n )\n hidden = torch.stack(next_rows).unsqueeze(0)\n\n liked_items = torch.from_numpy(\n prepared.inference_liked_items[output_rows].astype(np.int64)\n ).to(device)\n disliked_items = torch.from_numpy(\n prepared.inference_disliked_items[output_rows].astype(np.int64)\n ).to(device)\n queries = user_queries(\n model,\n hidden.squeeze(0),\n model._mean_context_channel(liked_items),\n model._mean_context_channel(disliked_items),\n torch.from_numpy(\n prepared.inference_user_features[output_rows].astype(np.float32)\n ).to(device),\n )\n result[torch.from_numpy(output_rows.astype(np.int64))] = queries.cpu()\n return result\n\n\[email protected]_mode()\ndef rank_candidates(\n model: base.SequentialTwoTower,\n prepared: SequenceData,\n eligibility: CandidateSet,\n config: Config,\n device: torch.device,\n ranking_batch_size: int,\n) -> tuple[pd.DataFrame, dict[str, float]]:\n model.eval()\n queries = inference_queries(model, prepared, config, device)\n item_parts: list[torch.Tensor] = []\n for start in range(0, len(prepared.item_ids), 2_048):\n rows = torch.arange(\n start,\n min(start + 2_048, len(prepared.item_ids)),\n dtype=torch.long,\n device=device,\n )\n item_parts.append(model.encode_items(rows).cpu())\n items = torch.cat(item_parts)\n rows: list[tuple[int, int, int]] = []\n for start in range(0, len(eligibility.uids), ranking_batch_size):\n end = min(start + ranking_batch_size, len(eligibility.uids))\n scores = queries[start:end] @ items.T\n eligible = np.unpackbits(\n eligibility.packed[start:end],\n axis=1,\n count=len(prepared.item_ids),\n bitorder=\"little\",\n ).astype(bool, copy=False)\n scores.masked_fill_(~torch.from_numpy(eligible), -torch.inf)\n selected = torch.topk(scores, k=base.TOP_N, dim=1).indices.numpy()\n for local_row, columns in enumerate(selected):\n uid = int(eligibility.uids[start + local_row])\n rows.extend(\n (uid, int(prepared.item_ids[column]), rank)\n for rank, column in enumerate(columns, start=1)\n )\n if end == len(eligibility.uids) or end % 1_024 == 0:\n log(f\"ranked users={end:,}/{len(eligibility.uids):,}\")\n query_norm = torch.linalg.vector_norm(queries, dim=1).numpy()\n item_norm = torch.linalg.vector_norm(items, dim=1).numpy()\n diagnostics = {\n \"query_norm_mean\": float(query_norm.mean()),\n \"query_norm_std\": float(query_norm.std()),\n \"item_norm_mean\": float(item_norm.mean()),\n \"item_norm_std\": float(item_norm.std()),\n }\n return pd.DataFrame(rows, columns=[\"uid\", \"item_id\", \"rank\"]), diagnostics\n\n\ndef write_ranking(frame: pd.DataFrame, output: Path) -> None:\n output.parent.mkdir(parents=True, exist_ok=True)\n if output.suffix == \".parquet\":\n frame.to_parquet(output, index=False, compression=\"zstd\")\n elif output.suffix == \".csv\":\n frame.to_csv(output, index=False)\n else:\n raise ValueError(\"--output must end in .parquet or .csv\")\n\n\ndef main() -> int:\n parser = argparse.ArgumentParser(description=__doc__)\n parser.add_argument(\n \"--stage\", choices=(\"validation\", \"submission\"), default=\"validation\"\n )\n parser.add_argument(\"--output\", type=Path, required=True)\n parser.add_argument(\"--public-dir\", type=Path, required=True)\n parser.add_argument(\"--prepared-cache\", type=Path)\n parser.add_argument(\"--device\", choices=(\"cpu\", \"mps\"), default=\"cpu\")\n parser.add_argument(\"--threads\", type=int, default=8)\n parser.add_argument(\"--bptt-steps\", type=int, default=Config.bptt_steps)\n parser.add_argument(\n \"--user-batch-size\", type=int, default=Config.user_batch_size\n )\n parser.add_argument(\"--epochs\", type=int, default=Config.epochs)\n parser.add_argument(\"--learning-rate\", type=float, default=Config.learning_rate)\n parser.add_argument(\n \"--loss-microbatch-size\",\n type=int,\n default=Config.loss_microbatch_size,\n )\n parser.add_argument(\n \"--negative-filter-scope\",\n choices=(\"full-training-history\", \"last-100\"),\n default=Config.negative_filter_scope,\n )\n parser.add_argument(\n \"--target-weighting\",\n choices=(\"discovery-graded\", \"organic-strong\", \"uniform\"),\n default=Config.target_weighting,\n )\n parser.add_argument(\n \"--ordinary-target-weight\",\n type=float,\n default=Config.ordinary_target_weight,\n )\n parser.add_argument(\n \"--organic-strong-target-weight\",\n type=float,\n default=Config.organic_strong_target_weight,\n )\n parser.add_argument(\n \"--organic-like-target-weight\",\n type=float,\n default=Config.organic_like_target_weight,\n )\n parser.add_argument(\n \"--discovery-horizon-days\",\n type=int,\n default=Config.discovery_horizon_days,\n )\n parser.add_argument(\n \"--validation-fraction\",\n type=float,\n default=Config.validation_fraction,\n )\n parser.add_argument(\n \"--max-train-targets-per-user\",\n type=int,\n default=Config.max_train_targets_per_user,\n help=\"retain this many recent pre-validation targets per user; 0 is uncapped\",\n )\n parser.add_argument(\n \"--max-validation-targets-per-user\",\n type=int,\n default=Config.max_validation_targets_per_user,\n help=\"retain this many final validation targets per user; 0 is uncapped\",\n )\n parser.add_argument(\n \"--early-stopping-patience\",\n type=int,\n default=Config.early_stopping_patience,\n )\n parser.add_argument(\n \"--early-stopping-min-delta\",\n type=float,\n default=Config.early_stopping_min_delta,\n )\n parser.add_argument(\"--checkpoint-dir\", type=Path)\n parser.add_argument(\"--resume-checkpoint\", type=Path)\n parser.add_argument(\"--seed\", type=int, default=Config.seed)\n parser.add_argument(\"--ranking-batch-size\", type=int, default=64)\n parser.add_argument(\"--prepare-only\", action=\"store_true\")\n args = parser.parse_args()\n\n if args.threads < 1 or args.bptt_steps < 2 or args.user_batch_size < 1:\n parser.error(\"threads, BPTT steps, and user batch size must be positive\")\n if (\n args.epochs < 1\n or args.learning_rate <= 0\n or args.loss_microbatch_size < 1\n or args.early_stopping_patience < 1\n or args.early_stopping_min_delta < 0\n or args.ordinary_target_weight <= 0\n or args.organic_strong_target_weight <= 0\n or args.organic_like_target_weight <= 0\n or args.discovery_horizon_days < 1\n or args.max_train_targets_per_user < 0\n or args.max_validation_targets_per_user < 0\n ):\n parser.error(\"training counts, learning rate, and patience are invalid\")\n if not 0.0 < args.validation_fraction < 1.0:\n parser.error(\"validation fraction must be strictly between zero and one\")\n public = args.public_dir.resolve()\n if not (public / \"data\" / \"candidate_catalog.parquet\").is_file():\n parser.error(f\"public task data not found under {public}\")\n config = Config(\n bptt_steps=args.bptt_steps,\n user_batch_size=args.user_batch_size,\n epochs=args.epochs,\n learning_rate=args.learning_rate,\n loss_microbatch_size=args.loss_microbatch_size,\n negative_filter_scope=args.negative_filter_scope,\n target_weighting=args.target_weighting,\n ordinary_target_weight=args.ordinary_target_weight,\n organic_strong_target_weight=args.organic_strong_target_weight,\n organic_like_target_weight=args.organic_like_target_weight,\n discovery_horizon_days=args.discovery_horizon_days,\n validation_fraction=args.validation_fraction,\n max_train_targets_per_user=args.max_train_targets_per_user,\n max_validation_targets_per_user=args.max_validation_targets_per_user,\n early_stopping_patience=args.early_stopping_patience,\n early_stopping_min_delta=args.early_stopping_min_delta,\n seed=args.seed,\n )\n torch.set_num_threads(args.threads)\n torch.set_num_interop_threads(1)\n seed_everything(config.seed)\n device = torch.device(args.device)\n if device.type == \"mps\" and not torch.backends.mps.is_available():\n parser.error(\"MPS was requested but is unavailable\")\n\n started = time.monotonic()\n prepared, eligibility = load_or_prepare(\n public,\n args.stage,\n config,\n args.prepared_cache.resolve() if args.prepared_cache else None,\n )\n target_partition = chronological_target_partition(\n prepared,\n config.validation_fraction,\n config.max_train_targets_per_user,\n config.max_validation_targets_per_user,\n )\n if args.prepare_only:\n print(\n json.dumps(\n {\n \"stage\": args.stage,\n \"sequence_users\": len(prepared.user_ids),\n \"sequence_tokens\": len(prepared.sequence_items),\n \"next_item_targets\": prepared.target_count,\n \"selected_training_targets\": int(\n target_partition.training_counts.sum()\n ),\n \"selected_validation_targets\": int(\n target_partition.validation_counts.sum()\n ),\n \"discarded_old_targets\": int(\n target_partition.window_starts.sum()\n ),\n \"target_users\": len(eligibility.uids),\n \"runtime_seconds\": round(time.monotonic() - started, 3),\n \"config\": asdict(config),\n },\n indent=2,\n )\n )\n return 0\n\n model, training_history = train_model(\n public,\n args.stage,\n"
}
}