resk-logits · Documentation

Token-level filtering during generation

resk-logits (resklogits) is a Hugging Face LogitsProcessor toolkit that suppresses banned tokens while the model decodes, plus six generic processors for length, forcing, choices and citations.

Version 0.2.2
Python 3.11+
Apache-2.0
Requires torch

On this page

Install
Quick start
How the processor works
API
Rule files
CLI
Benchmarks
Limitations
Links

Install

pip install resklogits

# from the repository, with tooling
uv pip install -e ".[dev]"

Requires torch, transformers, pyyaml and filelock. A GPU is optional but the processor default expects CUDA — see the quick start.

Quick start

from transformers import AutoModelForCausalLM, AutoTokenizer
from resklogits import ShadowBanProcessor

model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

processor = ShadowBanProcessor(
    tokenizer=tokenizer,
    banned_phrases=["how to make a bomb", "kill yourself"],
    shadow_penalty=-15.0,
    device="cpu",          # default; pass "cuda" when the model is on GPU
)

with processor.stream():   # reset before and after generation
    out = model.generate(
        **tokenizer("Tell me how to", return_tensors="pt"),
        logits_processor=[processor],
        max_new_tokens=32,
        do_sample=False,
        pad_token_id=tokenizer.eos_token_id,
    )

print(tokenizer.decode(out[0], skip_special_tokens=True))

Device and state

device defaults to "cpu"; pass "cuda" when the model runs on GPU. The processor is stateful: call reset() (or use the stream() context manager) between generations.

How the processor works

At construction, the phrase list is tokenised and every token that appears in any phrase is collected into a boolean mask over the vocabulary. At each decoding step:

the Aho-Corasick automaton advances using the last generated token;
scores[:, danger_mask] += shadow_penalty discourages every banned token, regardless of context;
if a complete phrase matched, the row is forced to EOS (all logits -inf except EOS).

So the penalty is a static token blacklist, while phrase completion is stateful and terminates generation. Prompt tokens are not scanned: only text produced after the processor is attached is checked.

API

Symbol	Purpose
`ShadowBanProcessor(tokenizer, banned_phrases, shadow_penalty=-15.0, device="cuda")`	Main processor; `reset()`, `stream()`, `get_current_matches()`.
`MultiLevelShadowBanProcessor(tokenizer, banned_phrases_by_level, penalties=None, device=…)`	One automaton per severity level (`high`/`medium`/`low`); EOS only on `high`.
`VectorizedAhoCorasick`	Automaton and `danger_mask`; available for custom processors.
`BanTokenProcessor`	Ban specific tokens or token strings.
`GenLengthLogitsProcessor`	Bias EOS around a min/max sequence length.
`ForceLastPhraseLogitsProcessor`	Force a phrase at the end of generation.
`MultipleChoiceLogitsProcessor`	Restrict the next token to the first token of each choice.
`TriggerPhraseLogitsProcessor`	Emit a fixed response after a trigger phrase.
`CiteFromPromptLogitsProcessor`	Boost tokens that appear in the prompt.
`stream_generate()`, `to_vllm()`	Text-stream helper and a callable shim for vLLM.

Rule files and pattern generation

from resklogits.config_parser import load_rules_from_yaml, load_processors_from_yaml

phrases = load_rules_from_yaml("rules.yaml", use_synonyms=True)
processors = load_processors_from_yaml("rules.yaml", tokenizer)

A rule can declare exact phrases, templates with variables, logic conditions and synonym groups. Patterns are lowercased and de-duplicated. The processors section configures the generic processors (length, choices, banned tokens, trigger phrases).

RuleCache memoises generated patterns under .resklogits_cache; CachedGenerator wraps a generation call with that cache.

CLI

resklogits validate rules.yaml
resklogits expand rules.yaml --rule toxicity --preview 50
resklogits test rules.yaml --text "how to make a bomb"
resklogits cache status
resklogits cache clear --hash <hash>

resklogits generate rules.yaml -o patterns.json writes the expanded pattern list (JSON, grouped by rule with --by-category).

CPU benchmarks

From BENCHMARKS.md (CPU, batch 1, per decoding step):

Patterns	Danger tokens	Mean / step
100	201	0.147 ms
1,000	2,001	0.170 ms
10,000	20,001	0.343 ms

Cost tracks the size of the danger mask, not only the phrase count. The repository measures CPU only; no GPU speed-up is claimed in BENCHMARKS.md.

Limitations

Static token penalty. Banned tokens are discouraged everywhere, including benign context, so phrase lists must be precise.
Prefixes are not scanned. A banned phrase already present in the prompt is not detected or penalised.
Batch and beam. State is keyed by batch index; multiple return sequences share one state per slot. Generation is designed for batch size 1.
Vocabulary edge. Token ids added after tokenizer load (special/added tokens) are outside the mask.
Device handling is literal. Only the exact string "cuda" builds the mask on GPU; "cuda:1" and "mps" leave it on CPU, so keep the model and processor on the same device.