perplexity.go

// Perplexity-guided selective token dropping (LLMLingua-style) — public API.
//
// Paper: "LLMLingua" — Jiang et al., Microsoft, 2023 (https://arxiv.org/abs/2310.05736)
//
// This file exposes an opt-in, pluggable compression path that drops the
// lowest-importance tokens from text, mirroring LLMLingua's perplexity-based
// pruning: tokens a language model finds predictable (low surprise) carry less
// information and are dropped first, while surprising / information-dense tokens
// are retained.
//
// The scoring is abstracted behind PerplexityScorer so callers can plug in a
// real model-backed scorer without changing the drop logic. The DEFAULT scorer
// (HeuristicPerplexityScorer) is a zero-dependency stand-in: it approximates
// token importance from local rarity + structural cues. It is a HEURISTIC, not
// true model perplexity — use a model-backed implementation when fidelity
// matters. tok ships an experimental OllamaScorer behind the experimental_ollama
// build tag.
//
// Usage:
//
//	out, stats := tok.CompressPerplexityGuided(text, tok.NewHeuristicPerplexityScorer(), 0.4)
//	out, stats := tok.Compress(text, tok.WithPerplexityGuided(scorer, 0.4))
package tok

import (
	"context"
	"math"
	"sort"
	"strings"
	"unicode"

	"github.com/GrayCodeAI/tok/internal/core"
)

// PerplexityScorer assigns an importance score to each token in a sequence.
// Higher scores mean more important / more surprising tokens that should be
// preserved; lower scores are dropped first. The returned slice MUST have the
// same length as tokens (one score per token); implementations that cannot
// honor that contract should return an error.
//
// The interface deliberately mirrors a model's per-token perplexity output so a
// real LM-backed scorer can be swapped in for the default heuristic without
// touching the drop logic.
type PerplexityScorer interface {
	Score(ctx context.Context, tokens []string) ([]float64, error)
}

// ScorerFunc adapts an ordinary function to the PerplexityScorer interface.
type ScorerFunc func(ctx context.Context, tokens []string) ([]float64, error)

// Score implements PerplexityScorer.
func (f ScorerFunc) Score(ctx context.Context, tokens []string) ([]float64, error) {
	return f(ctx, tokens)
}

// HeuristicPerplexityScorer is the DEFAULT zero-dependency scorer. It is a
// HEURISTIC stand-in for true model perplexity: it estimates token importance
// from purely local statistics, with no model calls, keeping tok dependency-light.
//
// Score(token) blends three signals:
//   - Inverse document frequency / rarity: tokens that occur rarely within the
//     text are more surprising and score higher (the perplexity proxy).
//   - Structural importance: identifiers, code-like tokens, numbers and
//     punctuation that carry syntactic weight are boosted.
//   - Length: longer tokens tend to be more specific and informative.
//
// This is NOT a substitute for a real language model. Plug in your own
// PerplexityScorer, or tok's experimental OllamaScorer when built with
// -tags experimental_ollama, when you need genuine perplexity.
type HeuristicPerplexityScorer struct {
	// RarityWeight scales the inverse-frequency (surprise) signal.
	RarityWeight float64
	// StructureWeight scales the structural-importance signal.
	StructureWeight float64
	// LengthWeight scales the token-length signal.
	LengthWeight float64
}

// NewHeuristicPerplexityScorer returns the default heuristic scorer with tuned
// weights. It is the zero-dependency default used by WithPerplexityGuided when
// no scorer is supplied.
func NewHeuristicPerplexityScorer() *HeuristicPerplexityScorer {
	return &HeuristicPerplexityScorer{
		RarityWeight:    1.0,
		StructureWeight: 1.0,
		LengthWeight:    0.5,
	}
}

// Score implements PerplexityScorer using local rarity + structural cues.
// It never returns an error: the heuristic is always defined for any input.
func (h *HeuristicPerplexityScorer) Score(_ context.Context, tokens []string) ([]float64, error) {
	scores := make([]float64, len(tokens))
	if len(tokens) == 0 {
		return scores, nil
	}

	// Document frequency within this text (case-folded for word tokens).
	freq := make(map[string]int, len(tokens))
	for _, t := range tokens {
		freq[strings.ToLower(t)]++
	}
	n := float64(len(tokens))

	for i, t := range tokens {
		// Rarity / inverse document frequency: rarer tokens are more surprising.
		// idf = log(N / df); a token appearing once scores near log(N).
		df := float64(freq[strings.ToLower(t)])
		rarity := math.Log((n + 1.0) / df)

		// Structural importance: protected/identifier/numeric/punctuation tokens.
		structure := structuralImportance(t)

		// Length: longer tokens are usually more specific.
		length := math.Log(float64(len([]rune(t))) + 1.0)

		scores[i] = h.RarityWeight*rarity +
			h.StructureWeight*structure +
			h.LengthWeight*length
	}
	return scores, nil
}

// structuralImportance returns a boost for tokens that are syntactically or
// semantically load-bearing: identifiers, code symbols, numbers, and
// sentence-terminating punctuation. Plain lowercase stop-words score ~0.
func structuralImportance(token string) float64 {
	if token == "" {
		return 0
	}
	if isProtectedToken(token) {
		// Protected symbols (sentence boundaries, code structure) are never the
		// cheapest thing to drop.
		return 3.0
	}

	var boost float64
	if looksLikeIdentifier(token) {
		boost += 2.0
	}
	if isNumericToken(token) {
		boost += 1.5
	}

	// Mixed-case / camelCase / PascalCase identifiers are highly specific.
	hasUpper, hasLower := false, false
	for _, r := range token {
		if unicode.IsUpper(r) {
			hasUpper = true
		}
		if unicode.IsLower(r) {
			hasLower = true
		}
	}
	if hasUpper && hasLower && len([]rune(token)) > 3 {
		boost += 1.0
	}
	return boost
}

// looksLikeIdentifier reports whether a token resembles a code identifier or
// path: contains underscores, dots between words, parens, sigils, or :: scope.
func looksLikeIdentifier(token string) bool {
	if strings.ContainsAny(token, "_()") || strings.Contains(token, "::") {
		return true
	}
	if strings.HasPrefix(token, "$") || strings.HasPrefix(token, "@") {
		return true
	}
	if strings.Contains(token, ".") && !strings.HasPrefix(token, ".") && !isNumericToken(token) {
		return true
	}
	return false
}

// isNumericToken reports whether a token is a number (allowing sign / decimal).
func isNumericToken(token string) bool {
	seenDigit := false
	for _, r := range token {
		switch {
		case r >= '0' && r <= '9':
			seenDigit = true
		case r == '.' || r == '-' || r == '+' || r == ',':
			// allowed separators
		default:
			return false
		}
	}
	return seenDigit
}

// isProtectedToken reports whether a token must never be dropped because it
// defines a sentence boundary or structural delimiter. Dropping these would
// corrupt the shape of the text.
func isProtectedToken(token string) bool {
	if token == "" {
		return false
	}
	// Pure punctuation/structure tokens.
	allPunct := true
	for _, r := range token {
		if !unicode.IsPunct(r) && !unicode.IsSymbol(r) {
			allPunct = false
			break
		}
	}
	if allPunct {
		return true
	}
	// Tokens ending in sentence-terminating punctuation carry boundary info.
	last := token[len(token)-1]
	return last == '.' || last == '!' || last == '?' || last == ':' || last == ';'
}

// WithPerplexityGuided enables an opt-in compression layer that drops the
// lowest-importance tokens (per scorer) up to dropRatio of the token count,
// while preserving sentence boundaries and protected symbols.
//
// dropRatio is clamped to [0, 0.9]; values <= 0 disable the layer. If scorer is
// nil, the default HeuristicPerplexityScorer is used.
//
// This runs as a final stage after the standard pipeline (it operates on the
// pipeline output), so it composes with tiers, modes, and budgets.
func WithPerplexityGuided(scorer PerplexityScorer, dropRatio float64) Option {
	return optFunc(func(c *config) {
		if scorer == nil {
			scorer = NewHeuristicPerplexityScorer()
		}
		c.perplexityScorer = scorer
		c.perplexityDropRatio = clampDropRatio(dropRatio)
		c.perplexityEnabled = dropRatio > 0
	})
}

// CompressPerplexityGuided applies ONLY perplexity-guided token dropping to text
// (no other pipeline layers). It is a self-contained entry point for callers who
// want the LLMLingua-style drop in isolation.
//
// If scorer is nil, the default HeuristicPerplexityScorer is used. dropRatio is
// clamped to [0, 0.9].
func CompressPerplexityGuided(text string, scorer PerplexityScorer, dropRatio float64) (string, Stats) {
	if text == "" {
		return "", Stats{}
	}
	if scorer == nil {
		scorer = NewHeuristicPerplexityScorer()
	}
	out := perplexityDrop(context.Background(), text, scorer, clampDropRatio(dropRatio))

	orig := core.EstimateTokens(text)
	final := core.EstimateTokens(out)
	saved := orig - final
	if saved < 0 {
		saved = 0
	}
	stats := Stats{
		OriginalTokens: orig,
		FinalTokens:    final,
		TokensSaved:    saved,
	}
	if orig > 0 {
		stats.ReductionPercent = float64(saved) / float64(orig) * 100
	}
	return out, stats
}

// perplexityDrop is the core LLMLingua-style drop. It tokenizes text on
// whitespace (preserving the original tokens), scores each token, then removes
// the lowest-scoring tokens up to dropRatio — but never drops protected tokens
// (sentence boundaries / structural symbols). On any scorer error or degenerate
// input it returns the input unchanged.
func perplexityDrop(ctx context.Context, text string, scorer PerplexityScorer, dropRatio float64) string {
	if dropRatio <= 0 {
		return text
	}
	tokens := strings.Fields(text)
	if len(tokens) < 4 {
		// Too little structure to prune safely.
		return text
	}

	scores, err := scorer.Score(ctx, tokens)
	if err != nil || len(scores) != len(tokens) {
		// Fail open: never corrupt output on a scorer failure.
		return text
	}

	// Candidates are droppable tokens (not protected), sorted by ascending score.
	type cand struct {
		idx   int
		score float64
	}
	candidates := make([]cand, 0, len(tokens))
	for i, t := range tokens {
		if isProtectedToken(t) {
			continue
		}
		candidates = append(candidates, cand{idx: i, score: scores[i]})
	}
	sort.SliceStable(candidates, func(i, j int) bool {
		return candidates[i].score < candidates[j].score
	})

	// Drop up to dropRatio of the FULL token count, taking from the lowest scores.
	dropTarget := int(math.Round(float64(len(tokens)) * dropRatio))
	if dropTarget > len(candidates) {
		dropTarget = len(candidates)
	}

	drop := make(map[int]struct{}, dropTarget)
	for i := 0; i < dropTarget; i++ {
		drop[candidates[i].idx] = struct{}{}
	}

	kept := make([]string, 0, len(tokens)-dropTarget)
	for i, t := range tokens {
		if _, dropped := drop[i]; dropped {
			continue
		}
		kept = append(kept, t)
	}
	return strings.Join(kept, " ")
}

// clampDropRatio bounds a drop ratio to a safe range. Negative -> 0, and we
// never drop more than 90% of tokens (beyond that output is unusable).
func clampDropRatio(r float64) float64 {
	if r < 0 {
		return 0
	}
	if r > 0.9 {
		return 0.9
	}
	return r
}