penalties.ts

/**
 * Repetition Penalties Implementation
 *
 * Ported from llama.cpp/src/llama-sampling.cpp (lines 1747-1778)
 * Exact formulas match native implementation for token-by-token parity
 *
 * References:
 * - Penalty apply: llama_sampler_penalties_apply (line 1747)
 * - Token tracking: llama_sampler_penalties_accept (line 1716)
 */

/**
 * Token history tracker with frequency map
 *
 * Matches llama.cpp's ring_buffer + unordered_map pattern
 * Maintains sliding window of last N tokens with O(1) operations
 */
export class TokenHistoryTracker {
  private ringBuffer: number[] = [];
  private tokenCount: Map<number, number> = new Map();
  private maxSize: number;

  constructor(penaltyLastN: number) {
    this.maxSize = Math.max(penaltyLastN, 0);
  }

  /**
   * Add token to history (matches llama_sampler_penalties_accept)
   */
  accept(token: number): void {
    if (this.maxSize === 0) {
      return;
    }

    // Increment count for new token
    this.tokenCount.set(token, (this.tokenCount.get(token) ?? 0) + 1);

    // If buffer is full, remove oldest token
    if (this.ringBuffer.length >= this.maxSize) {
      const oldest = this.ringBuffer.shift()!;

      // Decrement count for removed token
      const count = this.tokenCount.get(oldest)!;
      if (count === 1) {
        this.tokenCount.delete(oldest);
      } else {
        this.tokenCount.set(oldest, count - 1);
      }
    }

    // Add new token to back
    this.ringBuffer.push(token);
  }

  /**
   * Get occurrence count for a token in current window
   */
  getCount(token: number): number {
    return this.tokenCount.get(token) ?? 0;
  }

  /**
   * Check if token exists in current window
   */
  hasToken(token: number): boolean {
    return this.tokenCount.has(token);
  }

  /**
   * Reset history (matches llama_sampler_penalties_reset)
   */
  reset(): void {
    this.ringBuffer = [];
    this.tokenCount.clear();
  }

  /**
   * Get current window size
   */
  size(): number {
    return this.ringBuffer.length;
  }

  /**
   * Get array of unique token IDs in current window
   * Used for sparse penalty updates (O(H) instead of O(V))
   *
   * Performance: H ≈ 10-50 tokens vs V = 65,536
   */
  getUniqueTokens(): number[] {
    return Array.from(this.tokenCount.keys());
  }

  /**
   * Check if penalties would modify logits
   * Used to avoid unnecessary work
   */
  static hasPenalties(params: {
    repeat?: number;
    frequency?: number;
    presence?: number;
  }): boolean {
    return (
      (params.repeat !== undefined && params.repeat !== 1.0) ||
      (params.frequency !== undefined && params.frequency !== 0.0) ||
      (params.presence !== undefined && params.presence !== 0.0)
    );
  }

  /**
   * Compute penalty adjustment for a single token
   * Returns adjustment to subtract from logit (for freq/presence) or multiply factor (for repeat)
   *
   * This is the "virtual penalty" - O(1) lookup, no buffer mutation
   *
   * @param tokenId Token to check
   * @param baseLogit Original logit value
   * @param params Penalty parameters
   * @returns Adjusted logit after penalties
   */
  computeAdjustment(
    tokenId: number,
    baseLogit: number,
    params: {
      repeat?: number;
      frequency?: number;
      presence?: number;
    }
  ): number {
    const count = this.getCount(tokenId);
    if (count === 0) {
      return baseLogit; // No penalty if token not in history
    }

    let adjusted = baseLogit;

    // 1. Repetition penalty (multiplicative)
    if (params.repeat !== undefined && params.repeat !== 1.0) {
      if (adjusted <= 0) {
        adjusted *= params.repeat;
      } else {
        adjusted /= params.repeat;
      }
    }

    // 2. Frequency penalty (subtractive, scaled by count)
    if (params.frequency !== undefined && params.frequency !== 0.0) {
      adjusted -= count * params.frequency;
    }

    // 3. Presence penalty (subtractive, binary)
    if (params.presence !== undefined && params.presence !== 0.0) {
      adjusted -= params.presence;
    }

    return adjusted;
  }
}

/**
 * Apply repetition penalty (matches llama-sampling.cpp lines 1768-1772)
 *
 * CRITICAL: Different formula for negative vs positive logits
 * This matches the llama.cpp fix for the academic paper's bug
 *
 * OPTIMIZATION: Sparse iteration over H tokens (not V)
 * Performance: O(H) where H ≈ 10-50, not O(V) where V = 65,536
 *
 * @param logits Logits array (modified in-place)
 * @param tokenHistory Token history tracker
 * @param penalty Repetition penalty multiplier (typical: 1.0-1.5)
 */
export function applyRepetitionPenalty(
  logits: Float32Array,
  tokenHistory: TokenHistoryTracker,
  penalty: number
): void {
  if (penalty === 1.0) {
    return; // No-op
  }

  // Sparse iteration: only process tokens in history (H ≈ 10-50)
  const uniqueTokens = tokenHistory.getUniqueTokens();

  for (const tokenId of uniqueTokens) {
    // From llama.cpp comment (line 1766):
    // "The academic publication that described this technique actually
    // just only divided, but that would cause tokens with negative logits
    // to become more likely, which is obviously wrong.
    // This is common fix for this problem, which is to multiply by the
    // penalty instead of dividing."

    if (logits[tokenId] <= 0) {
      // Multiply for negative logits
      logits[tokenId] *= penalty;
    } else {
      // Divide for positive logits
      logits[tokenId] /= penalty;
    }
  }
}

/**
 * Apply frequency penalty (matches llama-sampling.cpp line 1774, first term)
 *
 * Penalty scales linearly with token occurrence count
 *
 * Formula: logit -= count * penalty_freq
 *
 * OPTIMIZATION: Sparse iteration over H tokens (not V)
 * Performance: O(H) where H ≈ 10-50, not O(V) where V = 65,536
 *
 * @param logits Logits array (modified in-place)
 * @param tokenHistory Token history tracker
 * @param penalty Frequency penalty (typical: 0.0-2.0)
 */
export function applyFrequencyPenalty(
  logits: Float32Array,
  tokenHistory: TokenHistoryTracker,
  penalty: number
): void {
  if (penalty === 0.0) {
    return; // No-op
  }

  // Sparse iteration: only process tokens in history (H ≈ 10-50)
  const uniqueTokens = tokenHistory.getUniqueTokens();

  for (const tokenId of uniqueTokens) {
    const count = tokenHistory.getCount(tokenId);
    // Exact formula from llama.cpp line 1774 (first term)
    logits[tokenId] -= count * penalty;
  }
}

/**
 * Apply presence penalty (matches llama-sampling.cpp line 1774, second term)
 *
 * Flat penalty if token appears at least once (binary, not scaled by count)
 *
 * Formula: logit -= (count > 0 ? 1 : 0) * penalty_present
 *
 * OPTIMIZATION: Sparse iteration over H tokens (not V)
 * Performance: O(H) where H ≈ 10-50, not O(V) where V = 65,536
 *
 * @param logits Logits array (modified in-place)
 * @param tokenHistory Token history tracker
 * @param penalty Presence penalty (typical: 0.0-1.0)
 */
export function applyPresencePenalty(
  logits: Float32Array,
  tokenHistory: TokenHistoryTracker,
  penalty: number
): void {
  if (penalty === 0.0) {
    return; // No-op
  }

  // Sparse iteration: only process tokens in history (H ≈ 10-50)
  const uniqueTokens = tokenHistory.getUniqueTokens();

  for (const tokenId of uniqueTokens) {
    // Exact formula from llama.cpp line 1774 (second term)
    // Note: All tokens in uniqueTokens have count > 0
    logits[tokenId] -= penalty;
  }
}

/**
 * Apply all penalties in correct order (matches llama-sampling.cpp line 1774)
 *
 * Order matches llama.cpp:
 * 1. Repetition penalty (multiply/divide based on logit sign)
 * 2. Frequency + Presence penalties (subtractive)
 *
 * Note: llama.cpp applies frequency and presence in a single pass (line 1774),
 * but the effect is identical to applying them separately
 *
 * @param logits Logits array (modified in-place)
 * @param tokenHistory Token history tracker
 * @param params Penalty parameters
 */
export function applyPenalties(
  logits: Float32Array,
  tokenHistory: TokenHistoryTracker,
  params: {
    repeat?: number;
    frequency?: number;
    presence?: number;
  }
): void {
  // Early return optimization (matches llama.cpp lines 1750-1752)
  if (
    tokenHistory.size() === 0 ||
    (params.repeat === 1.0 && params.frequency === 0.0 && params.presence === 0.0)
  ) {
    return;
  }

  // Apply in correct order (matches llama.cpp)
  if (params.repeat !== undefined && params.repeat !== 1.0) {
    applyRepetitionPenalty(logits, tokenHistory, params.repeat);
  }

  if (params.frequency !== undefined && params.frequency !== 0.0) {
    applyFrequencyPenalty(logits, tokenHistory, params.frequency);
  }

  if (params.presence !== undefined && params.presence !== 0.0) {
    applyPresencePenalty(logits, tokenHistory, params.presence);
  }
}