Top-K Logit Filteringeasy

Top-K Logit Filtering

easy Inside an LLM 15–20 min

NumPy allowed · plain Python earns 10 bonus points

Logits become a next-token distribution

Softmax turns a model score vector into probabilities.

Open lesson →

01prompt

02tokens

03logits

04probabilities

attention→next token

mechanism65%

layer20%

pattern10%

model5%

Softmax(logits) → probability distribution

Filter vocabulary logits to the top-K highest scoring tokens before sampling.

Task

Implement top_k_filter(logits, k, filter_value=-float('inf')). Retain only the $k$ tokens with the highest logits. Replace all other logits with filter_value. If multiple tokens tie at the cutoff, prioritize lower token indices. If $k ge ext{len}(logits)$, return an unmasked copy of logits. If $k le 0$, raise ValueError.

Function Signature

top_k_filter(logits: list[float], k: int, filter_value: float = -float('inf')) -> list[float]

Examples

Example 1: Top 2 tokens kept
Input: logits = [4, 1, 3, 0.5]
k = 2
Output: [4, float('-inf'), 3, float('-inf')]

Example 2: k >= vocab size returns copy
Input: logits = [2, 1]
k = 5
Output: [2, 1]

Example 3: Invalid k <= 0 error
Input: logits = [2, 1]
k = 0
Output: Raises ValueError

Constraints

k <= 0 must raise ValueError.
If k >= len(logits), return original logits copy without masking.
Ties at cutoff must be broken by smaller index.
A plain-Python solution earns 10 bonus points.

Python3Saved ✓

def top_k_filter(logits: list[float], k: int, filter_value: float = -float('inf')) -> list[float]:

"""Mask all logits outside the top-k highest scoring tokens.

Args:

logits: Raw token scores.

k: Number of highest scoring tokens to retain.

filter_value: Value assigned to masked tokens (default -inf).

Returns:

Filtered logits list.

"""

# Your implementation here

pass

Case 1: Top 2 tokens kept

Input: logits = [4, 1, 3, 0.5]
k = 2

Expected: [4, float('-inf'), 3, float('-inf')]

Case 2: k >= vocab size returns copy

Input: logits = [2, 1]
k = 5

Expected: [2, 1]

Case 3: Invalid k <= 0 error

Input: logits = [2, 1]
k = 0

Expected: Raises ValueError