| |
| |
| |
|
|
| |
| |
| """ |
| Segment-Bound Masking for Ranker Attention. |
| |
| Segment IDs: HISTORY=1, CANDIDATE=-1, PADDING=0 |
| Ranker mask: history attends to all, candidates attend only to self. |
| """ |
|
|
| from dataclasses import dataclass |
|
|
| import jax |
| import jax.numpy as jnp |
|
|
| HISTORY_SEGMENT_ID = 1 |
| CANDIDATE_SEGMENT_ID = -1 |
| PADDING_SEGMENT_ID = 0 |
|
|
|
|
| @dataclass(frozen=True) |
| class SegmentBounds: |
| history_lower: jax.Array |
| history_upper: jax.Array |
| candidate_lower: jax.Array |
| candidate_upper: jax.Array |
|
|
| @classmethod |
| def from_segment_ids(cls, segment_ids: jax.Array) -> "SegmentBounds": |
| segment_ids = jnp.asarray(segment_ids) |
| if segment_ids.ndim == 1: |
| segment_ids = segment_ids[None, :] |
| B, S = segment_ids.shape |
|
|
| history_mask = segment_ids == HISTORY_SEGMENT_ID |
| candidate_mask = segment_ids == CANDIDATE_SEGMENT_ID |
|
|
| def _bounds(mask): |
| any_mask = jnp.any(mask, axis=1) |
| first = jnp.argmax(mask, axis=1) |
| last = (S - 1) - jnp.argmax(mask[:, ::-1], axis=1) |
| lower = jnp.where(any_mask, first, 0) |
| upper = jnp.where(any_mask, last + 1, 0) |
| return lower.astype(jnp.int32), upper.astype(jnp.int32) |
|
|
| hl, hu = _bounds(history_mask) |
| cl, cu = _bounds(candidate_mask) |
| return cls(hl, hu, cl, cu) |
|
|
| def to_array(self) -> jax.Array: |
| return jnp.stack([self.history_lower, self.history_upper, |
| self.candidate_lower, self.candidate_upper], axis=1) |
|
|
|
|
| def ranker_mask(q_pos: jax.Array, kv_pos: jax.Array, bounds: SegmentBounds) -> jax.Array: |
| q_is_hist = (q_pos >= bounds.history_lower[..., None, None]) & \ |
| (q_pos < bounds.history_upper[..., None, None]) |
| q_is_cand = (q_pos >= bounds.candidate_lower[..., None, None]) & \ |
| (q_pos < bounds.candidate_upper[..., None, None]) |
| kv_is_hist = (kv_pos >= bounds.history_lower[..., None, None]) & \ |
| (kv_pos < bounds.history_upper[..., None, None]) |
| kv_is_cand = (kv_pos >= bounds.candidate_lower[..., None, None]) & \ |
| (kv_pos < bounds.candidate_upper[..., None, None]) |
|
|
| history_mask = kv_is_hist & (q_is_hist | q_is_cand) |
| candidate_self_mask = q_is_cand & kv_is_cand & (q_pos == kv_pos) |
| return history_mask | candidate_self_mask |
|
|
|
|
| def causal_ranker_mask(q_pos: jax.Array, kv_pos: jax.Array, bounds: SegmentBounds, |
| window_len: int = -1) -> jax.Array: |
| base = ranker_mask(q_pos, kv_pos, bounds) |
| causal = q_pos >= kv_pos |
| if window_len > 0: |
| causal = causal & (kv_pos > q_pos - window_len) |
| return base & causal |
|
|
