llama: add BertForTokenClassification support - #19725
Open
giladgd wants to merge 5 commits into
Open
Conversation
giladgd
requested review from
CISC,
danbev,
ggerganov,
ngxson and
pwilkin
as code owners
February 19, 2026 06:24
Contributor
Author
|
@CISC I saw you reviewed and added support for classification models in the past, so it'd be great if you could review this PR |
richiejp
added a commit
to richiejp/LocalAI
that referenced
this pull request
Jun 12, 2026
…k queue Score and TokenClassify previously bypassed llama.cpp's slot scheduler and drove llama_decode directly on the shared llama_context. That required static mutexes, a fatal conflict_guard tripwire that aborted the process if they ever overlapped the slot loop, and a config-validation rule forbidding score/token_classify from sharing a llama-cpp model config with chat/completion/embeddings. Both now ride the server task queue: - patch 0001 gains PR ggml-org/llama.cpp#19725's tools/server hunks (re-based to the pin): send_embedding treats TOKEN_CLS pooling as token-level, so embedding tasks return one raw n_cls logit row per token. TokenClassify submits its overlapping windows as embedding tasks and keeps log-softmax/Viterbi/span stitching in the handler. - new patch 0006 adds SERVER_TASK_TYPE_SCORE: the batch builder flags candidate positions for logits, a per-chunk harvest accumulates log_softmax(logits)[next_token] after every decode (a scored region can span multiple n_batch chunks, which also lifts the input cap to n_ctx), and send_score emits the summed result at DONE_PROMPT. Prompt-cache reuse is clamped to score_start-1 so the position predicting the first scored token is always re-decoded; beyond that, cross-candidate prompt-KV reuse comes free from the slot prompt cache, retiring the old re-decode-per-candidate perf TODO. server_n_outputs_max gains bounded headroom (+64) for score outputs and the batch fill enforces that budget per decode - without this the upstream n_outputs_max ceiling (sized for last-token-only logits) asserts, which also means the old direct-decode Score violated it at the current pin. - the conflict_guard machinery, both RPC mutexes, and the ModelConfig.Validate usecase-conflict blocks are deleted; a single llama-cpp config may now combine score/token_classify with chat/completion/embeddings, and the GGUF importer guard keeps its behavior (reserved models stay out of chat pickers) with the rationale updated. Verified locally against the CPU grpc-server build: TokenClassify spans match the known-good output with 4 concurrent calls returning identical results; score rankings are sane with exact per-token logprob counts, bitwise-identical single-chunk results across n_batch 512 vs 64, chunk- invariant (fp-noise) multi-chunk results including a 212-token candidate through the output-budget path; Score and Predict run concurrently with no abort; pooled embeddings still L2-normalize (dims 768, norm 1.0). All six patches apply at pin 7c158fbb with --fuzz=0, idempotently. Assisted-by: claude-code:claude-fable-5 [Claude Code]
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR adds token classification support for
BertForTokenClassificationandModernBertForTokenClassification, which allows me to use this model and also this one.I added a new pooling type for token-level classification since it makes
llama_get_embeddings_ithreturn an array with a length ofllama_model_n_cls_out(model)(instead ofllama_model_n_embd_out(model)).AI disclosure: I used Codex for assistance but wrote most of the code myself.