Tool response
1→#!/usr/bin/env python3
"""Train the Yambda windowed two-tower baseline.
Every catalog listen with at least 50 percent completion is a sequence token.
Within each selected chronological window, every adjacent target contributes a
loss. The fast default retains each user's most recent 128 training targets and
16 held-out validation targets; passing zero for either cap restores the
uncapped behavior for that side of the split. A 100-step truncated-BPTT block
computes all of its timestep losses and carries the detached GRU hidden state
10→into the next block.
The model is a pure two-tower retriever: no user-ID embedding, no output
normalization, raw dot-product retrieval, sample-weighted binary cross entropy,
and 24 uniform plus eight propensity-corrected unique in-batch negatives. Every
selected adjacent target retains a nonzero weight. Only files under the supplied
public directory are read.
"""
from __future__ import annotations
20→import argparse
import json
import random
import sys
import time
from dataclasses import asdict, dataclass
from pathlib import Path
import numpy as np
import pandas as pd
30→import torch
from torch import nn
from torch.nn import functional as F
import two_tower_model as base
CACHE_VERSION = 2
40→@dataclass(frozen=True)
class Config:
bptt_steps: int = 100
user_batch_size: int = 64
embedding_dim: int = 192
hidden_dim: int = 192
user_hidden_1: int = 512
user_hidden_2: int = 256
item_hidden_1: int = 512
item_hidden_2: int = 256
50→ dropout: float = 0.05
epochs: int = 1
learning_rate: float = 1.0e-3
weight_decay: float = 1.0e-5
gradient_clip: float = 2.0
loss_microbatch_size: int = 65_536
uniform_negatives: int = 24
inbatch_negatives: int = 8
inbatch_correction_power: float = 0.5
inbatch_correction_min: float = 0.25
60→ inbatch_correction_max: float = 4.0
negative_filter_scope: str = "full-training-history"
target_weighting: str = "organic-strong"
ordinary_target_weight: float = 0.25
organic_strong_target_weight: float = 1.0
organic_like_target_weight: float = 3.0
discovery_horizon_days: int = 7
validation_fraction: float = 0.10
max_train_targets_per_user: int = 128
max_validation_targets_per_user: int = 16
70→ early_stopping_patience: int = 1
early_stopping_min_delta: float = 1.0e-4
seed: int = 29
@dataclass(frozen=True)
class CandidateSet:
uids: np.ndarray
item_ids: np.ndarray
packed: np.ndarray
80→
@classmethod
def read(cls, path: Path) -> "CandidateSet":
with np.load(path, allow_pickle=False) as artifact:
result = cls(
uids=artifact["uids"].astype(np.int64, copy=True),
item_ids=artifact["item_ids"].astype(np.int64, copy=True),
packed=artifact["packed_eligible"].astype(np.uint8, copy=True),
)
expected = (len(result.uids), (len(result.item_ids) + 7) // 8)
90→ if result.packed.shape != expected:
raise ValueError(
f"invalid candidate artifact shape {result.packed.shape}; "
f"expected {expected}"
)
return result
@dataclass(frozen=True)
class SequenceData:
100→ item_ids: np.ndarray
item_features: np.ndarray
user_ids: np.ndarray
offsets: np.ndarray
sequence_timestamps: np.ndarray
sequence_items: np.ndarray
event_features: np.ndarray
activity_features: np.ndarray
state_ids: np.ndarray
liked_states: np.ndarray
110→ disliked_states: np.ndarray
state_counts: np.ndarray
inference_user_rows: np.ndarray
inference_liked_items: np.ndarray
inference_disliked_items: np.ndarray
inference_user_features: np.ndarray
@property
def lengths(self) -> np.ndarray:
return np.diff(self.offsets)
120→
@property
def target_count(self) -> int:
return int(np.maximum(self.lengths - 1, 0).sum())
@dataclass(frozen=True)
class TargetPartition:
"""A contiguous recent target window and its chronological BCE split."""
130→ window_starts: np.ndarray
training_counts: np.ndarray
validation_counts: np.ndarray
history_training_counts: np.ndarray
@property
def selected_count(self) -> int:
return int(self.training_counts.sum() + self.validation_counts.sum())
140→def log(message: str) -> None:
print(message, file=sys.stderr, flush=True)
def seed_everything(seed: int) -> None:
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
150→def model_config(config: Config) -> base.TowerConfig:
"""Map the dense trainer configuration onto the unchanged tower."""
return base.TowerConfig(
sequence_length=config.bptt_steps,
embedding_dim=config.embedding_dim,
hidden_dim=config.hidden_dim,
user_hidden_1=config.user_hidden_1,
user_hidden_2=config.user_hidden_2,
item_hidden_1=config.item_hidden_1,
item_hidden_2=config.item_hidden_2,
160→ dropout=config.dropout,
epochs=config.epochs,
learning_rate=config.learning_rate,
weight_decay=config.weight_decay,
gradient_clip=config.gradient_clip,
uniform_negatives=config.uniform_negatives,
inbatch_negatives=config.inbatch_negatives,
inbatch_correction_power=config.inbatch_correction_power,
inbatch_correction_min=config.inbatch_correction_min,
inbatch_correction_max=config.inbatch_correction_max,
170→ seed=config.seed,
)
def append_state_snapshot(
active_likes: dict[int, int],
active_dislikes: dict[int, int],
n_items: int,
liked_states: list[np.ndarray],
disliked_states: list[np.ndarray],
180→ state_counts: list[tuple[int, int]],
) -> int:
liked_states.append(base.pack_active_items(active_likes, base.LIKE_ITEMS, n_items))
disliked_states.append(
base.pack_active_items(active_dislikes, base.DISLIKE_ITEMS, n_items)
)
state_counts.append((len(active_likes), len(active_dislikes)))
return len(liked_states) - 1
190→def causal_activity_features(
timestamps: np.ndarray,
items: np.ndarray,
organic: np.ndarray,
completion: np.ndarray,
) -> np.ndarray:
"""Compute the first four reviewed user features at every event cutoff."""
length = len(items)
result = np.zeros((length, 4), dtype=np.float16)
if not length:
200→ return result
left_rows = np.searchsorted(
timestamps,
timestamps - base.PROFILE_LOOKBACK_DAYS * base.DAY,
side="left",
)
rows = np.arange(length, dtype=np.int64)
counts = rows - left_rows + 1
organic_sum = np.cumsum(organic.astype(np.float64))
completion_sum = np.cumsum(np.minimum(completion, 100).astype(np.float64))
210→ before_organic = np.where(
left_rows > 0, organic_sum[np.maximum(left_rows - 1, 0)], 0.0
)
before_completion = np.where(
left_rows > 0, completion_sum[np.maximum(left_rows - 1, 0)], 0.0
)
unique_counts = np.empty(length, dtype=np.int32)
frequencies: dict[int, int] = {}
left = 0
220→ for right, raw_item in enumerate(items):
item = int(raw_item)
frequencies[item] = frequencies.get(item, 0) + 1
required_left = int(left_rows[right])
while left < required_left:
old_item = int(items[left])
remaining = frequencies[old_item] - 1
if remaining:
frequencies[old_item] = remaining
else:
230→ del frequencies[old_item]
left += 1
unique_counts[right] = len(frequencies)
result[:, 0] = (np.log1p(counts) / 8.0).astype(np.float16)
result[:, 1] = (np.log1p(unique_counts) / 6.0).astype(np.float16)
result[:, 2] = (
(organic_sum - before_organic) / np.maximum(counts, 1)
).astype(np.float16)
result[:, 3] = (
240→ (completion_sum - before_completion)
/ np.maximum(counts, 1)
/ 100.0
).astype(np.float16)
return result
def build_sequence_parts(
listens: base.ListenEvents,
explicit: base.ExplicitEvents,
250→ n_items: int,
liked_states: list[np.ndarray],
disliked_states: list[np.ndarray],
state_counts: list[tuple[int, int]],
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray] | None:
qualifying = listens.completion_pct >= base.LISTEN_POSITIVE_THRESHOLD
timestamps = listens.timestamps[qualifying].astype(np.int64, copy=False)
items = listens.items[qualifying].astype(np.uint16, copy=False)
organic = listens.organic[qualifying].astype(np.float32, copy=False)
completion = listens.completion_pct[qualifying].astype(np.float32, copy=False)
260→ if not len(items):
return None
event_features = np.zeros((len(items), 3), dtype=np.float16)
event_features[:, 0] = organic.astype(np.float16)
event_features[:, 1] = (
np.minimum(completion, 100.0) / 100.0
).astype(np.float16)
if len(items) > 1:
gaps = np.maximum(0, timestamps[1:] - timestamps[:-1])
270→ event_features[1:, 2] = np.exp(
-gaps.astype(np.float64) / (60.0 * base.DAY)
).astype(np.float16)
activity = causal_activity_features(
timestamps, items, organic, completion
)
states = np.zeros(len(items), dtype=np.uint32)
active_likes: dict[int, int] = {}
active_dislikes: dict[int, int] = {}
280→ explicit_position = 0
current_state = 0
for row, timestamp in enumerate(timestamps):
changed = False
while (
explicit_position < len(explicit.timestamps)
and int(explicit.timestamps[explicit_position]) < int(timestamp)
):
item = int(explicit.items[explicit_position])
event_timestamp = int(explicit.timestamps[explicit_position])
290→ destination = (
active_likes
if int(explicit.channel[explicit_position]) == 0
else active_dislikes
)
if int(explicit.active[explicit_position]) == 1:
destination[item] = event_timestamp
else:
destination.pop(item, None)
explicit_position += 1
300→ changed = True
if changed:
current_state = append_state_snapshot(
active_likes,
active_dislikes,
n_items,
liked_states,
disliked_states,
state_counts,
)
310→ states[row] = np.uint32(current_state)
return timestamps, items.copy(), event_features, activity, states
def final_user_features(
timestamps: np.ndarray,
items: np.ndarray,
event_features: np.ndarray,
history_end: int,
like_count: int,
320→ dislike_count: int,
) -> np.ndarray:
start = int(
np.searchsorted(
timestamps,
history_end - base.PROFILE_LOOKBACK_DAYS * base.DAY,
side="left",
)
)
result = np.zeros(6, dtype=np.float32)
330→ recent_items = items[start:]
if len(recent_items):
result[0] = np.log1p(len(recent_items)) / 8.0
result[1] = np.log1p(len(np.unique(recent_items))) / 6.0
result[2] = float(np.mean(event_features[start:, 0]))
result[3] = float(np.mean(event_features[start:, 1]))
result[4] = np.log1p(like_count) / 6.0
result[5] = np.log1p(dislike_count) / 4.0
return result
340→
def prepare_data(
public: Path, stage: str, config: Config
) -> tuple[SequenceData, CandidateSet]:
data = public / "data"
catalog = pd.read_parquet(data / "candidate_catalog.parquet").sort_values(
"popularity_rank"
)
item_ids = catalog.item_id.to_numpy(dtype=np.int64, copy=True)
item_index = pd.Index(item_ids)
350→ n_items = len(item_ids)
if n_items >= np.iinfo(np.uint16).max:
raise RuntimeError("uint16 item storage is too small for this catalog")
if stage == "validation":
prefixes = ("train",)
history_end = base.TRAIN_END
candidate_name = "validation_eligible_candidates.npz"
state_name = "validation_preference_state.parquet"
else:
360→ prefixes = ("train", "val")
history_end = base.VALIDATION_END
candidate_name = "test_eligible_candidates.npz"
state_name = "test_preference_state.parquet"
eligibility = CandidateSet.read(data / candidate_name)
if not np.array_equal(eligibility.item_ids, item_ids):
raise RuntimeError("catalog and candidate artifact item order differ")
explicit_by_user = base.load_explicit_events(data, prefixes, item_index)
inference_likes, inference_dislikes, like_counts, dislike_counts = (
370→ base.load_inference_state(
data / state_name,
item_index,
eligibility.uids,
n_items,
)
)
inference_rows = {int(uid): row for row, uid in enumerate(eligibility.uids)}
inference_user_features = np.zeros((len(eligibility.uids), 6), dtype=np.float32)
inference_sequence_rows = np.full(len(eligibility.uids), -1, dtype=np.int32)
380→
# Snapshot zero is the shared empty explicit state.
liked_states = [np.full(base.LIKE_ITEMS, n_items, dtype=np.uint16)]
disliked_states = [
np.full(base.DISLIKE_ITEMS, n_items, dtype=np.uint16)
]
state_counts: list[tuple[int, int]] = [(0, 0)]
user_ids: list[int] = []
offsets = [0]
timestamp_parts: list[np.ndarray] = []
390→ item_parts: list[np.ndarray] = []
event_parts: list[np.ndarray] = []
activity_parts: list[np.ndarray] = []
state_parts: list[np.ndarray] = []
listen_paths = [data / f"{prefix}_listens.parquet" for prefix in prefixes]
started = time.monotonic()
for seen_users, (uid, listens) in enumerate(
base.merge_listen_streams(listen_paths, item_index), start=1
):
400→ parts = build_sequence_parts(
listens,
explicit_by_user.get(uid, base.empty_explicit_events()),
n_items,
liked_states,
disliked_states,
state_counts,
)
if parts is None:
continue
410→ sequence_row = len(user_ids)
user_ids.append(uid)
timestamp_parts.append(parts[0].astype(np.uint32, copy=False))
item_parts.append(parts[1])
event_parts.append(parts[2])
activity_parts.append(parts[3])
state_parts.append(parts[4])
offsets.append(offsets[-1] + len(parts[1]))
inference_row = inference_rows.get(uid)
if inference_row is not None:
420→ inference_sequence_rows[inference_row] = sequence_row
inference_user_features[inference_row] = final_user_features(
parts[0],
parts[1],
parts[2],
history_end,
int(like_counts[inference_row]),
int(dislike_counts[inference_row]),
)
if seen_users % 1_000 == 0:
430→ log(
f"prepared users={seen_users:,} tokens={offsets[-1]:,} "
f"states={len(liked_states):,} "
f"elapsed={time.monotonic() - started:.1f}s"
)
missing = np.flatnonzero(inference_sequence_rows < 0)
if len(missing):
inference_user_features[missing, 4] = np.log1p(like_counts[missing]) / 6.0
inference_user_features[missing, 5] = (
440→ np.log1p(dislike_counts[missing]) / 4.0
)
log(f"target users without qualifying catalog sequence={len(missing):,}")
prepared = SequenceData(
item_ids=item_ids,
item_features=base.build_item_features(catalog),
user_ids=np.asarray(user_ids, dtype=np.int64),
offsets=np.asarray(offsets, dtype=np.int64),
sequence_timestamps=np.concatenate(timestamp_parts),
450→ sequence_items=np.concatenate(item_parts),
event_features=np.concatenate(event_parts),
activity_features=np.concatenate(activity_parts),
state_ids=np.concatenate(state_parts),
liked_states=np.stack(liked_states),
disliked_states=np.stack(disliked_states),
state_counts=np.asarray(state_counts, dtype=np.int32),
inference_user_rows=inference_sequence_rows,
inference_liked_items=inference_likes,
inference_disliked_items=inference_dislikes,
460→ inference_user_features=inference_user_features,
)
log(
f"prepared stage={stage} sequence_users={len(prepared.user_ids):,} "
f"tokens={len(prepared.sequence_items):,} "
f"next_item_targets={prepared.target_count:,} "
f"states={len(prepared.liked_states):,}"
)
return prepared, eligibility
470→
def cache_metadata(stage: str, config: Config) -> dict[str, object]:
return {
"cache_version": CACHE_VERSION,
"stage": stage,
"positive_listen_threshold": base.LISTEN_POSITIVE_THRESHOLD,
"profile_lookback_days": base.PROFILE_LOOKBACK_DAYS,
"event_time_feature": "exp-negative-gap-over-60-days",
}
480→
def save_cache(
path: Path, prepared: SequenceData, stage: str, config: Config
) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
np.savez(
path,
metadata=np.asarray(json.dumps(cache_metadata(stage, config), sort_keys=True)),
**prepared.__dict__,
)
490→ log(f"saved public-derived full-sequence cache: {path}")
def load_cache(path: Path, stage: str, config: Config) -> SequenceData:
with np.load(path, allow_pickle=False) as artifact:
actual = json.loads(str(artifact["metadata"].item()))
expected = cache_metadata(stage, config)
if actual != expected:
raise RuntimeError(f"sequence cache differs: {actual} != {expected}")
values = {
500→ field: artifact[field].copy()
for field in SequenceData.__dataclass_fields__
}
prepared = SequenceData(**values)
log(
f"loaded full-sequence cache: tokens={len(prepared.sequence_items):,} "
f"targets={prepared.target_count:,}"
)
return prepared
510→
def load_or_prepare(
public: Path,
stage: str,
config: Config,
cache: Path | None,
) -> tuple[SequenceData, CandidateSet]:
candidate_name = (
"validation_eligible_candidates.npz"
if stage == "validation"
520→ else "test_eligible_candidates.npz"
)
eligibility = CandidateSet.read(public / "data" / candidate_name)
if cache is not None and cache.is_file():
prepared = load_cache(cache, stage, config)
if not np.array_equal(prepared.item_ids, eligibility.item_ids):
raise RuntimeError("cached catalog differs from candidate artifact")
return prepared, eligibility
prepared, built_eligibility = prepare_data(public, stage, config)
if cache is not None:
530→ save_cache(cache, prepared, stage, config)
return prepared, built_eligibility
def build_item_correction(
prepared: SequenceData,
config: Config,
target_starts: np.ndarray,
target_counts: np.ndarray,
) -> torch.Tensor:
540→ """Estimate in-batch propensities from the selected training targets."""
counts = np.zeros(len(prepared.item_ids), dtype=np.float64)
for user_row, count in enumerate(target_counts):
if not count:
continue
offset = int(prepared.offsets[user_row])
target_start = int(target_starts[user_row])
counts += np.bincount(
prepared.sequence_items[
offset
550→ + target_start
+ 1 : offset
+ target_start
+ int(count)
+ 1
].astype(np.int64),
minlength=len(prepared.item_ids),
)
total = float(counts.sum())
probabilities = (counts + 1.0) / (total + len(counts))
560→ raw = np.power(
(1.0 / len(counts)) / probabilities,
config.inbatch_correction_power,
)
normalization = float(np.sum(counts * raw) / max(total, 1.0))
correction = np.clip(
raw / max(normalization, 1.0e-12),
config.inbatch_correction_min,
config.inbatch_correction_max,
).astype(np.float32)
570→ mean_on_targets = float(np.sum(counts * correction) / max(total, 1.0))
log(
f"inbatch correction: min={correction.min():.3f} "
f"mean_on_targets={mean_on_targets:.3f} max={correction.max():.3f}"
)
return torch.from_numpy(correction)
def build_known_positive_filter(
prepared: SequenceData, training_counts: np.ndarray
580→) -> np.ndarray:
"""Pack each user's complete public-training positive-item set."""
width = (len(prepared.item_ids) + 7) // 8
packed = np.zeros((len(prepared.user_ids), width), dtype=np.uint8)
retained = 0
for user_row, target_count in enumerate(training_counts):
if not target_count:
continue
offset = int(prepared.offsets[user_row])
# T adjacent targets consume sequence positions 0 through T. The
590→ # chronological BCE suffix starts at T + 1 and is intentionally absent.
items = np.unique(
prepared.sequence_items[
offset : offset + int(target_count) + 1
].astype(np.int64)
)
byte_columns = items >> 3
bit_values = np.left_shift(
np.uint8(1), (items & 7).astype(np.uint8)
)
600→ np.bitwise_or.at(packed[user_row], byte_columns, bit_values)
retained += len(items)
log(
f"full-history false-negative filter: users={len(prepared.user_ids):,} "
f"unique_user_items={retained:,} bytes={packed.nbytes:,}"
)
return packed
def invalid_full_history_negatives(
610→ negatives: np.ndarray,
targets: np.ndarray,
user_rows: np.ndarray,
known_positive_packed: np.ndarray,
liked_items: np.ndarray,
) -> np.ndarray:
"""Return row-level invalid negatives using the complete training prefix."""
byte_columns = negatives >> 3
bit_columns = negatives & 7
known = known_positive_packed[user_rows[:, None], byte_columns]
620→ invalid = ((known >> bit_columns) & 1).astype(bool)
invalid |= negatives == targets[:, None]
invalid |= (negatives[:, :, None] == liked_items[:, None, :]).any(axis=2)
for column in range(1, negatives.shape[1]):
invalid[:, column] |= (
negatives[:, column, None] == negatives[:, :column]
).any(axis=1)
return invalid
630→def draw_full_history_filtered_negatives(
targets: torch.Tensor,
user_rows: np.ndarray,
liked_items: torch.Tensor,
known_positive_packed: np.ndarray,
n_items: int,
config: Config,
item_correction: torch.Tensor,
generator: torch.Generator,
device: torch.device,
640→) -> tuple[torch.Tensor, torch.Tensor]:
"""Draw mixed negatives and reject all known public-training positives."""
targets_cpu = targets.detach().cpu().numpy().astype(np.int64, copy=False)
liked_cpu = liked_items.detach().cpu().numpy().astype(np.int64, copy=False)
batch_size = len(targets_cpu)
uniform = torch.randint(
n_items,
(batch_size, config.uniform_negatives),
generator=generator,
device="cpu",
650→ )
inbatch = torch.empty((batch_size, 0), dtype=torch.long)
corrected_count = 0
if config.inbatch_negatives:
unique_targets = torch.unique(targets.detach().cpu())
corrected_count = min(
config.inbatch_negatives, max(len(unique_targets) - 1, 0)
)
if corrected_count:
priorities = torch.rand(
660→ (batch_size, len(unique_targets)), generator=generator
)
priorities.masked_fill_(
unique_targets[None, :].eq(
torch.from_numpy(targets_cpu)[:, None]
),
-1.0,
)
columns = torch.topk(
priorities, k=corrected_count, dim=1
670→ ).indices
inbatch = unique_targets[columns]
if corrected_count < config.inbatch_negatives:
filler = torch.randint(
n_items,
(batch_size, config.inbatch_negatives - corrected_count),
generator=generator,
device="cpu",
)
inbatch = torch.cat([inbatch, filler], dim=1)
680→
negatives = torch.cat([uniform, inbatch], dim=1).numpy()
correction = np.ones_like(negatives, dtype=np.float32)
if corrected_count:
correction[
:, config.uniform_negatives : config.uniform_negatives + corrected_count
] = item_correction[inbatch[:, :corrected_count]].numpy()
for _ in range(32):
invalid = invalid_full_history_negatives(
negatives,
690→ targets_cpu,
user_rows,
known_positive_packed,
liked_cpu,
)
if not np.any(invalid):
return (
torch.from_numpy(negatives).to(device),
torch.from_numpy(correction).to(device),
)
700→ replacements = torch.randint(
n_items,
(int(invalid.sum()),),
generator=generator,
device="cpu",
).numpy()
negatives[invalid] = replacements
# Rejected in-batch items become ordinary uniform draws.
correction[invalid] = 1.0
invalid = invalid_full_history_negatives(
710→ negatives,
targets_cpu,
user_rows,
known_positive_packed,
liked_cpu,
)
if np.any(invalid):
raise RuntimeError("full-history negative resampling failed")
return (
torch.from_numpy(negatives).to(device),
720→ torch.from_numpy(correction).to(device),
)
def draw_training_negatives(
positive_items: torch.Tensor,
known_history: torch.Tensor,
liked_items: torch.Tensor,
user_rows: np.ndarray,
known_positive_packed: np.ndarray | None,
730→ prepared: SequenceData,
config: Config,
item_correction: torch.Tensor,
generator: torch.Generator,
device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor]:
if config.negative_filter_scope == "full-training-history":
if known_positive_packed is None:
raise RuntimeError("full-history negative filter was not built")
return draw_full_history_filtered_negatives(
740→ positive_items,
user_rows,
liked_items,
known_positive_packed,
len(prepared.item_ids),
config,
item_correction,
generator,
device,
)
750→ return base.draw_filtered_mixed_negatives(
positive_items,
known_history,
liked_items,
len(prepared.item_ids),
config.uniform_negatives,
config.inbatch_negatives,
item_correction,
generator,
device,
760→ )
def event_sequence(
model: base.SequentialTwoTower,
items: torch.Tensor,
features: torch.Tensor,
hidden: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
inputs = torch.cat([model.context_embedding(items), features], dim=-1)
770→ projected = model.event_projection(inputs)
return model.sequence_encoder(projected, hidden)
def pool_explicit_states(
model: base.SequentialTwoTower,
prepared: SequenceData,
state_ids: np.ndarray,
device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
780→ unique_states, inverse = np.unique(state_ids, return_inverse=True)
liked_items = torch.from_numpy(
prepared.liked_states[unique_states].astype(np.int64)
).to(device)
disliked_items = torch.from_numpy(
prepared.disliked_states[unique_states].astype(np.int64)
).to(device)
liked_vectors = model._mean_context_channel(liked_items)
disliked_vectors = model._mean_context_channel(disliked_items)
inverse_rows = torch.from_numpy(inverse.astype(np.int64)).to(device)
790→ counts = prepared.state_counts[state_ids]
return (
liked_vectors[inverse_rows],
disliked_vectors[inverse_rows],
torch.from_numpy(counts.astype(np.float32)).to(device),
)
def user_queries(
model: base.SequentialTwoTower,
800→ sequence_vectors: torch.Tensor,
liked_vectors: torch.Tensor,
disliked_vectors: torch.Tensor,
user_features: torch.Tensor,
) -> torch.Tensor:
return model.user_tower(
torch.cat(
[sequence_vectors, liked_vectors, disliked_vectors, user_features],
dim=1,
)
810→ )
def candidate_vectors(
model: base.SequentialTwoTower,
positives: torch.Tensor,
negatives: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
all_items = torch.cat([positives, negatives.reshape(-1)])
unique_items, inverse = torch.unique(all_items, return_inverse=True)
820→ encoded = model.encode_items(unique_items)
positive_vectors = encoded[inverse[: len(positives)]]
negative_vectors = encoded[inverse[len(positives) :]].reshape(
len(positives), negatives.shape[1], -1
)
return positive_vectors, negative_vectors
def history_windows(
sequence: np.ndarray,
830→ start: int,
count: int,
width: int,
padding_item: int,
) -> np.ndarray:
base_start = max(0, start - width + 1)
segment = sequence[base_start : start + count]
padded = np.concatenate(
[np.full(width - 1, padding_item, dtype=np.uint16), segment]
)
840→ windows = np.lib.stride_tricks.sliding_window_view(padded, width)
offset = start - base_start
return windows[offset : offset + count].copy()
def corrected_negative_loss(
logits: torch.Tensor,
correction: torch.Tensor,
config: Config,
) -> torch.Tensor:
850→ raw = F.softplus(logits)
total = config.uniform_negatives + config.inbatch_negatives
result = torch.zeros(len(logits), device=logits.device)
if config.uniform_negatives:
result += config.uniform_negatives * raw[
:, : config.uniform_negatives
].mean(dim=1)
if config.inbatch_negatives:
selected = slice(config.uniform_negatives, total)
weights = correction[:, selected]
860→ result += config.inbatch_negatives * (
(raw[:, selected] * weights).sum(dim=1)
/ weights.sum(dim=1).clamp_min(1.0e-6)
)
return result / total
def retained_target_window_starts(
total_targets: np.ndarray,
validation_counts: np.ndarray,
870→ training_counts: np.ndarray,
) -> np.ndarray:
"""Return the discarded prefix before the retained chronological suffix."""
training_region_ends = total_targets - validation_counts
window_starts = training_region_ends - training_counts
if (
np.any(window_starts < 0)
or np.any(training_counts < 0)
or np.any(validation_counts < 0)
):
880→ raise RuntimeError("chronological target partition produced negative counts")
if not np.array_equal(
window_starts + training_counts + validation_counts, total_targets
):
raise RuntimeError("chronological target window is not a final contiguous suffix")
return window_starts
def chronological_target_partition(
prepared: SequenceData,
890→ validation_fraction: float,
max_train_targets_per_user: int,
max_validation_targets_per_user: int,
) -> TargetPartition:
"""Select a recent contiguous window with a final out-of-time BCE suffix."""
totals = np.maximum(prepared.lengths - 1, 0).astype(np.int64)
validation = np.floor(totals * validation_fraction).astype(np.int64)
validation[(totals >= 2) & (validation == 0)] = 1
if max_validation_targets_per_user:
validation = np.minimum(validation, max_validation_targets_per_user)
900→
history_training = totals - validation
training = history_training.copy()
if max_train_targets_per_user:
training = np.minimum(training, max_train_targets_per_user)
window_starts = retained_target_window_starts(totals, validation, training)
return TargetPartition(
window_starts=window_starts,
training_counts=training,
validation_counts=validation,
910→ history_training_counts=history_training,
)
def load_organic_like_outcomes(
public: Path, stage: str, item_ids: np.ndarray
) -> dict[int, dict[int, np.ndarray]]:
"""Load mapped organic-like timestamps from public fitting history."""
prefixes = ("train",) if stage == "validation" else ("train", "val")
frames = [
920→ pd.read_parquet(
public / "data" / f"{prefix}_likes.parquet",
columns=["uid", "timestamp", "item_id", "is_organic"],
)
for prefix in prefixes
]
likes = pd.concat(frames, ignore_index=True)
likes = likes[likes.is_organic == 1].copy()
mapped = pd.Index(item_ids).get_indexer(
likes.item_id.to_numpy(dtype=np.int64)
930→ )
likes = likes.loc[mapped >= 0, ["uid", "timestamp"]].copy()
likes["item"] = mapped[mapped >= 0].astype(np.int32, copy=False)
likes = likes.sort_values(["uid", "item", "timestamp"], kind="stable")
nested: dict[int, dict[int, list[int]]] = {}
for raw_uid, raw_timestamp, raw_item in likes.itertuples(
index=False, name=None
):
nested.setdefault(int(raw_uid), {}).setdefault(int(raw_item), []).append(
int(raw_timestamp)
940→ )
result = {
uid: {
item: np.asarray(timestamps, dtype=np.int64)
for item, timestamps in items.items()
}
for uid, items in nested.items()
}
log(
f"loaded organic-like outcomes: rows={len(likes):,} "
950→ f"users={len(result):,}"
)
return result
def build_target_weights(
public: Path,
stage: str,
prepared: SequenceData,
training_counts: np.ndarray,
960→ config: Config,
) -> np.ndarray:
"""Build nonzero per-token weights without dropping adjacent targets."""
if config.target_weighting == "uniform":
weights = np.ones(len(prepared.sequence_items), dtype=np.float32)
else:
weights = np.full(
len(prepared.sequence_items),
config.ordinary_target_weight,
dtype=np.float32,
970→ )
organic_strong = (
(prepared.event_features[:, 0] > 0.5)
& (prepared.event_features[:, 1] >= 0.8)
)
weights[organic_strong] = config.organic_strong_target_weight
if config.target_weighting == "discovery-graded":
like_outcomes = load_organic_like_outcomes(
public, stage, prepared.item_ids
)
980→ history_end = (
base.TRAIN_END if stage == "validation" else base.VALIDATION_END
)
horizon = config.discovery_horizon_days * base.DAY
grade_three = 0
for user_row, uid in enumerate(prepared.user_ids):
start = int(prepared.offsets[user_row])
end = int(prepared.offsets[user_row + 1])
if end - start < 2:
continue
990→ last_training_target = start + int(training_counts[user_row])
first_validation_target = last_training_target + 1
split_timestamp = (
int(prepared.sequence_timestamps[first_validation_target])
if first_validation_target < end
else history_end
)
user_likes = like_outcomes.get(int(uid), {})
seen = {int(prepared.sequence_items[start])}
for position in range(start + 1, end):
1000→ item = int(prepared.sequence_items[position])