c608649552* feat(studio): run chats in parallel in the Chat tab
New Chat used to cancel whatever the current conversation was generating.
It now leaves it running, like switching to the Train or Export tab: the
sidebar shows which chats are still going, and Stop is per conversation.
Plain `unsloth studio` launched llama-server with one decode slot, so the
admission queue serialised every chat regardless of what the UI did. Both
entry points now default to the same slot count as `unsloth studio run`.
A model swap still ends every running chat, since they all decode on one
llama-server. /load and /unload now refuse with 409 and name those chats
unless the caller passes force_cancel_active, and the UI asks first.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix(studio): scope the composer tool badge to its own conversation
The green "Running Python: ..." badge above the composer read a single
global store value, so one chat's tool call showed above every other
chat's composer, including a brand-new empty one. Its elapsed counter
also restarted at 0 on every thread switch, and a run ending anywhere
cleared the badge everywhere.
Key the status by thread and store the moment it started, so each
conversation shows only its own tool call and the counter resumes rather
than restarts. Also adds a test that every conversation gets its own
tool sandbox directory, which parallel tool calls depend on.
* Fix stalled tool calls while awaiting approval for PR #7455
Three problems, all from the approval prompt behaving as though only one
chat could ever run.
Arguments were not streamed for a gated call, so the chat stayed blank for
as long as the model took to write the payload, which for a large file is
minutes. Nothing runs before the decision either way, and the code is what
is being approved, so python and terminal now stream their card while
gated. render_html stays suppressed: its card renders the payload.
The status read "Running ..." with a climbing timer while the call had not
started. It now reports that it is waiting for approval, then switches to
running once allowed.
The admission lease was held across the wait, so four unanswered prompts
held all four decode slots and no other chat could start while llama-server
sat idle. A parked run keeps its lease but no longer counts against
capacity.
Measured with four prompts left open: every gated call streamed its code,
none reported running, and a fresh chat answered in 0.4s where it
previously waited 290s and never did.
* Fix duplicated and truncated tool cards for PR #7455
A gated tool call rendered two cards: the provisional one that streams the
arguments, plus a second one keyed by the approval id. Only the second ever
got its tool_end, so the first spun "Running" for the rest of the chat.
Reuse the open part when the approval prompt arrives.
The terminal card also showed nothing but a 60-char trigger label, so a long
heredoc read as no progress at all. It now renders the command the same way
the Python card renders its script, and neither is capped at 10k chars.
Both cells moved inside the collapsible, so one chevron hides the code with
the output and Copy / Download exist only while the card is open. A card
parked on the prompt says so instead of counting up "Running".
* Fix review findings on the parallel-chat gate for PR #7455
Backend:
- /unload rechecks active generations under the lifecycle gate, like /load,
and lets its 409 through the catch-all instead of rewriting it as a 500.
- /load gates only once _load_model_impl has decided this is a real reload,
so an Apply on the already-loaded model no longer refuses, and the retry it
asks for no longer cancels every chat before returning already_loaded.
- The direct /v1/responses stream registers in the cancel registry, so a
non-forced unload can no longer tear llama-server down under it.
- run_server defaults to the same slot count as the CLI. colab.py calls it
without the argument, so Colab was still serialising every chat.
Frontend:
- Cancelling a backgrounded chat aborts its own request rather than only
posting a cancel id, which is the only thing that ends an external-provider
or audio run.
- The model-swap dialog counts local runs only, and falls back to the backend
when this tab's map is empty, so a reload or a second tab still gets asked.
- Context usage and the diffusion canvas are scoped to the chat that produced
them; a compare row reads activity from its member threads.
Tests:
- The extracted-source cancel harnesses supply the active-generations module,
which the tracked-cancel class now depends on.
* Fix the swap confirmation scope and cancel timing for PR #7455
A forced load cancelled every chat before the model identifier, GPU selection,
training coexistence and download checks had run, so a load that then failed
those checks stopped the chats and replaced nothing. The refusal still happens
early, but the destructive cancel now sits immediately before the teardown it
is paying for, and rechecks under the gate like /unload does.
The swap dialog only reconciled with the backend when this tab looked idle, so
one local chat was enough to hide a second tab's runs. Confirming then sent
force_cancel_active, which cancels every backend run, including the ones the
dialog never mentioned. The backend snapshot is now merged in every time, so
the dialog names what will actually stop. External-provider runs are never
registered there, so the union stays local-only.
Also drops the active-generations docstring claim about restoring sidebar
spinners, which nothing consumes.
* Defer destructive cancels and track every local stream for PR #7455
/unload cancelled the running chats before it had resolved that it unloads
anything. A stale model_path, which a second tab produces routinely, killed
every chat and then no-opped, leaving the resident model up. It now refuses
early and cancels only at each teardown, matching /load.
The swap dialog also stopped every chat locally the moment the user confirmed,
which threw away the two-phase backend behaviour: a load that then failed
identifier resolution, GPU validation or the training guard had already
truncated the replies. The backend now owns the cancel.
Three local streams decoded on llama-server without registering, so a
non-forced unload counted zero generations and tore the server down mid
response: /v1/completions streaming, and the plain and server-tool Anthropic
streams, the first of which is the default /v1/messages path. Note this makes
a non-forced load return 409 during those runs rather than draining quietly,
the same trade the /v1/responses fix made.
The safetensors tool loop still announced a gated call as running while it
waited on a human; only the GGUF loop had been fixed. A source-level parity
test now pins both.
Also drops stopAllChatThreads, which has no callers left.
* Studio: close three load/unload gate races found in review
Re-check the in-flight load guard after the stop-running-chats confirm.
The confirm always GETs active-generations before its zero-running
early-out, so the guard no longer sits atomically ahead of the
reservation and two picks in that window both reached performLoad over
the same refs. ejectModel had the same shape and gets the same re-check.
Reject a sidecar swap immediately before the forced cancel in both load
branches. The previous check was back at the top of preflight, so an
install reserving during identifier resolution, the tier probe, the
training guard or the download check made the post-drain recheck 409 a
load whose chats had already been stopped.
Enter the Anthropic passthrough's cancel tracker inside its body
generator. It was entered eagerly and returned through
_sse_streaming_response, which sets no unstarted_cleanup, so a response
whose body never started left the run registered forever and 409'd every
later non-forced load and unload.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Trim comments across the files this PR touches
Tightens the comments and doc blocks in the backend, CLI, tests and frontend
files changed by this PR: collapses multi-line explanations to a single line
where they still read clearly, and drops the ones the code already says.
No code changes, verified by an AST comparison against the previous commit.
* Studio: defer the destructive cancel and close two gate gaps
Move the forced cancel behind every check that can still reject a swap.
The drain now runs first with the runs it is about to cancel discounted,
so it waits only for inference the cancel cannot end, then the sidecar
check decides, then the cancel fires, then a second drain lets those runs
unwind before teardown. A sidecar install reserving during the drain no
longer 409s a load whose chats have already been stopped.
Track the non-streaming /v1/completions proxy. It was the last local
decode path missing from active_generations, so an unload, which runs no
drain, tore llama-server down under it and force_cancel_active could not
signal it. It now uses the same tracked cancel event and dedicated client
as the OpenAI pass-through.
Skip the client's preliminary unload while chats are generating and let
/load evict at its own post-preflight point instead. Forwarding
force_cancel_active there truncated replies before identifier
resolution, the GPU and training guards and the download check had run.
Keep per-thread context usage so returning to a chat whose background run
finished restores its bar instead of leaving it blank until the next turn.
Make the running-flag clear run-specific. Every run without a resolved
thread id shares the "__default" key, so concurrent compare panes could
clear each other's flag and strand a live stop handle.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: register the embeddings proxy with the swap gate
/v1/embeddings proxied straight through the pooled client with no tracked
cancel event, so it never appeared in active_generations. /unload runs no
idle drain, so a concurrent non-forced unload counted zero generations and
killed llama-server mid-request, and force_cancel_active had no event to
signal. Mirrors the completions proxy: tracked event, dedicated unpooled
client closed by a cancel/disconnect watcher, unregister in a nested
finally so a close failure cannot leave a phantom generation behind.
* Trim comments on the newest changes in this PR
Comments only, no code changes: shorten the ones added by the load-gate
ordering, embeddings and per-thread usage work down to the same density as
the rest of the diff.
* Studio: register the legacy generate stream with the swap gate
/generate/stream built a cancel event but never entered the tracker, so it
was invisible to active_generations. Being in the keep-warm middleware's
inference suffixes only covers /load, which drains; /unload does not, so a
non-forced unload passed the 409 gate and then blocked on the standard
backend's generation lock, and a forced swap had no event to signal.
Registered inside the body generator under a nested finally so a teardown
failure cannot skip the unregister.
The AST contract test asserted the cleanup finally by overwriting its flag
per Try node, so a nested try made the last one win. Accumulate instead,
which is what the existence claim meant.
* Studio: three more swap-gate gaps found in review
Register /audio/generate with the gate. TTS holds the model for the whole
request and /unload runs no drain, so unregistered a non-forced swap counted
zero generations and tore the model down mid-generation; the orchestrator
path only waits 15s for the generation lock, which real TTS exceeds. No
cancel keys: no backend takes a cancel_event for audio, so the event has no
observer and a forced swap still cannot interrupt audio already in flight.
Thread the tracked cancel event into the /v1/responses admission wait. It
was the only admission caller passing None, so a queued run could not be
reached by cancel_all() and a plain /inference/cancel could not stop it at
all. Same omission fixed at the upstream send there and on /v1/completions.
Let an unforced unload of a stale model path reach the no-op check. Before
this PR that request returned 200 and did nothing; the new gate refused it
with 409 for a request that reaches no teardown branch. Gate both refusal
passes on the disjunction of the route's own teardown conditions, including
not is_loaded, so a mid-load GGUF still refuses.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: register the remaining non-streaming decode paths
stream defaults to false on all three of these, so they are the ordinary
shape of their routes, and each holds a local backend for the whole
request. /unload runs no idle drain, so with no registry entry a non-forced
swap counted zero generations and tore the backend down mid-request instead
of returning 409, and a forced one had no event to signal.
Non-streaming /v1/messages: all three helpers ran with an empty registry,
since only the streaming siblings were tracked. Registered at the call site
because the pass-through takes no cancel_event of its own, and with no
cancel keys, matching those siblings.
Non-streaming standard chat and audio-input chat: the trackers in this route
sit inside their `if payload.stream:` arms, so neither else branch was
covered. The GGUF sibling already registers its own non-streaming branch.
Each exit is in a finally on the branch's existing try, so the except arms
are covered too: a leaked entry 409s every later swap until restart.
* Studio: tighten the swap-gate comments
Comment-only pass over the newest swap-gate registrations: collapse the multi-line rationales in /unload, the legacy generate stream, audio generation and the non-streaming chat branches, and the matching test preambles, to the shortest form that still carries the reason. No code changes.
* Studio: stop the reselect dialog promising a stop that never happens
Picking an external provider leaves the local model resident and stops the
status poll mirroring it, so reselecting that model showed the stop-chats
dialog, and /load then answered already_loaded ahead of its cancel hook.
Confirmed with the live backend: the same pick with force_cancel_active set
still returned already_loaded and the chat kept streaming. Not stopping
those chats is right, since the load never interrupts them, so remove the
prompt rather than honour it. Blanket-skipping is unsafe, because the same
id and variant with one sampling setting changed is a real reload and 409s,
so the branch only fires when a status fetch confirms the resident
checkpoint and variant match, and then adopts it without calling /load.
Redact native model paths from the active-generations response. Registering
/generate/stream recorded backend.active_model_name verbatim, which is an
absolute path for a native local model, and this route is the only place
that serialises it. Redacting at the response covers every tracker rather
than the one that surfaced it.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: keep hydrated context usage in the per-thread map
The history loader restores a saved conversation's usage through
setContextUsage only, and it runs once per mount, so switching away and
back left the bar blank for a hydrated chat even after the per-thread map
landed. setContextUsage now writes the value through to the visible
thread's own entry and clears that entry when passed null, which covers
both hydration call sites and any future writer.
* Studio: unblock load cancellation and share unresolved thread keys
Run the two stop-loading fast paths ahead of the unload route's pre-gate
refusal. _unload_may_evict returns True for exactly the model being
cancelled, so the refusal was blocking the branch that cancels a load which
has replaced nothing and can interrupt no chat. The client made that
unrecoverable: cancelLoading sends the unload without force, drops the
result, and its abort never reaches /load, which takes no signal, so the
load ran on and could later cancel those chats and swap the model. Nothing
else is exempted; an unload that would tear down a serving model matches
neither fast path and still 409s. The comment claiming the client lets that
409 surface is corrected, since it discards it.
Hold every owner behind a shared thread key. Runs with no resolved thread id
share "__default" (concurrent compare panes, since startCompare clears
activeThreadId), so a single owner slot let a second run replace the first's
token and then delete the shared entry while it was still generating, and
the server-cancel map lost the older handle the same way. Both now hold a
list, the running and local flags survive until the last owner clears, and
stopChatThread stops every handle under the key.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: carry a confirmed swap into the sidecar install, key restored usage by thread
Picking a model that needs a newer transformers while chats generate raised the
"stop N chats" prompt, but the answer never reached the install that runs before
the load: /install-latest-transformers refused on those same chats and took no
force flag, so Retry hit the same 409 and nothing in the flow stopped them.
Carry force_cancel_active through the consent dialog into the installer. Only
the pre-gate fast path is skipped: the recheck under the lifecycle gate still
has to pass, so an unconfirmed caller is refused as before. The cancel runs last
inside the gate, after every check that can still reject the install, and the
drain behind it is bounded since it holds the gate and the sidecar reservation.
Also key restored context usage by the thread the loader read. history.load()
captures remoteId before two awaited round trips, so a switch inside that window
filed one thread's usage under another and setActiveThreadId kept re-applying it.
Preserve sibling owners when a run key is cleared without an owner: the image
rejection gate now uses its own token, and the reducer leaves owned runs alone.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: bound the post-cancel drains, and make cancellation reach the paths that ignored it
A forced swap cancels the chats it interrupts, then waits for them to unwind.
That wait had no deadline while holding the lifecycle gate, and TTS on the
subprocess backend observes no cancel event at all, so one audio generation
could pin every load, unload and new request for its whole duration. Bound both
post-cancel drains. Pre-cancel drains stay unbounded: the swap can still be
refused there, so shortening them would weaken what they protect.
/unload had the opposite problem and no drain at all, cancelling and tearing
down on the next line, which turned a clean stream end into a dropped
connection. Give it the same bounded wait, gated on the cancel having cancelled
something so an idle Eject pays nothing.
Make the cancel actually land where it can. GGUF TTS now takes a cancel_event
and a watcher closes its client to break the blocking POST. The Anthropic
non-streaming pass-through did the same thing the completions and embeddings
paths used to: register with the gate, then run both POSTs on the pooled client
that cannot be closed. It now uses a per-request client like they do.
Also: park and unpark the admission queue the reservation actually holds, since
queues are keyed by base_url and a reload mints a new port; key tool output by
remoteId on both sides, so the first turn of a New Chat stops writing under one
key and reading another; and give tool status a run owner, so a finishing run
cannot blank the badge a concurrent one is still showing.
Clamp --parallel to 1 on a llama-server without --kv-unified. The new default of
4 would otherwise split -c four ways on such a build, quartering the context
window for a feature it cannot serve.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: scope a chat's Stop to its own generation, and clear the way before a confirmed install
Safetensors generation is serialized on _gen_lock and the worker has a single
cancel event, so a chat still queued on that lock owns no generation. Its Stop
handler called reset_generation_state() anyway, which set the shared event and
ended whichever conversation was actually running. Parallel chats is what makes
that reachable.
_generate_inner now records its cancel_event as the current holder once it takes
the lock, and reset_generation_state drops a reset from anyone else. Every route
call site passes its own request event. A reset with no event stays global, so
unload and model switch cannot leave a generation alive, and a reset while
nothing runs still resets, so an error path before generation is not a no-op.
The other two backends take the argument too, or the standard one raises
TypeError on every cancel.
The sidecar install had the mirror of the /load ordering problem: it cancelled
the chats first and drained second, so an unrelated counted request the cancel
cannot reach (a count_tokens, say) was still there for the recheck, which then
refused an install that had already stopped every chat for nothing. Drain the
unreachable remainder first, discounting the registered chats, then cancel.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: close the windows the previous round's fixes left open
Three follow-ups, two of them holes in the fixes just before them.
The worker claim went in after _send_cmd, so the command was already running
unclaimed and a queued chat's Stop in that window still reset it. Claim first,
with the send inside the same try, so a failed send releases it too.
Tool status kept one entry per key with an owner. That stops a foreign clear but
not an overwrite: under the shared unresolved-thread key the second run replaced
the first's entry, and its own clear then removed the only one while the first
tool was still running. Keep per-run entries and render the newest.
/unload gated its drain on having cancelled something, so a request that passed
the keep-warm middleware but had not reached its tracker yet was invisible to it
and the teardown landed on an already-admitted request. Drain on the middleware
count instead, which covers that window as well as the cancelled runs, then
re-cancel whatever registered while waiting. Bounded, not a refusal: an unload is
deliberate, and on expiry it proceeds exactly as before.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Trim the parallel-chats comments to their reasons
Compress the multi-line rationales added by this branch into shorter forms and drop
restatements of the code below them. The reasons behind the drain bounds, the deferred
cancel, the per-request generation ownership and the thread-scoped tool and usage keys
are kept, just said in fewer lines.
* Studio: own the worker per generation, and make a resumed chat requeue for its slot
Ownership was a single lock holder, so dispatched runs (compare mode bypasses
_gen_lock by design) never claimed it and the guard fell straight through to the
global reset: a Stop on one of them ended its siblings. Track the generations
actually running instead, claimed before the send and released in the same
finally on both paths. A reset still proceeds when nothing is running, so an
error path ahead of generation is not swallowed.
park() hands the freed slot to a waiter, so a chat resuming from a tool approval
could take it back while that waiter was still decoding, putting two holders on
a one-slot server and sending the resumed tool loop past the admission limit.
unpark_async waits for room; the plain unpark stays for a holder tearing down,
which will not decode again.
Audio only observed its cancel event on a forced swap. An explicit Stop just
aborts the fetch, and this route has no cancel id, so llama-server ran on to the
request timeout after the chat reported it stopped. Watch the disconnect.
Also read tool status by remoteId, matching the key the adapter writes and the
fix already made for tool output, and stop an unresolved run from writing its
usage into whichever conversation the user moved to.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: let only the generation the worker is running speak for it, and hold a slot for a resuming chat
The ownership list recorded admission, but the subprocess runs generations one
at a time, so a dispatched request queued behind another counted as an owner and
its Stop signalled the shared cancel event, ending the request that was actually
running. Keep admission for release bookkeeping and gate ownership on execution
instead, promoted when the worker first answers that request. Nothing executing
still permits a reset, so an error path ahead of generation is not swallowed.
The worker has one cancel event and no per-request cancellation, so this decides
who may pull the lever rather than making the lever per-request.
A resuming chat also polled for a slot it could never see: release() grants to
the next waiter under the same lock, so later arrivals overtook an approved chat
indefinitely. A pending unpark now reserves the next slot and they queue behind
it.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: cover the prefill window, and keep a first turn's tool output readable
Gating worker ownership on execution left the interval between the send and the
first response uncovered: nothing is executing then, and the empty case admitted
anyone, so a queued chat's Stop still ended the one in prefill. Split the empty
case. Nothing claimed at all still permits a reset, so an error path ahead of
generation is not swallowed; claimed but unanswered resolves to the oldest
claim, which is what a FIFO command queue is working on.
Putting both sides of the tool-output scope on remoteId left the first turn of a
New Chat writing under the unresolved scope for its whole life while the readers
recomputed the moment the autosave assigned an id, so the card blanked mid-run.
The readers now fall back to the unresolved scope, which only an unpersisted
first turn can occupy.
* Studio: order the parked approvals, and tie a worker claim to its enqueue
The reservation added for admission fairness was a bare count, so every approved
holder counted against every other: park two chats, approve both, and once the
last decoder released, nothing could ever satisfy the check again. That is a
deadlock where the problem it fixed was only unfairness. Make it a FIFO ticket
so a pending unpark blocks the ones behind it and no others.
_owns_worker reads claim order to decide which request the worker is prefilling,
which only holds if claiming and enqueuing cannot interleave. Hold one lock
across both on the dispatched and the locked path.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: follow a first turn's run when its thread gets an id, and key the denoising canvas per chat
A run started before its thread existed filed every handle under "__default". Nothing
moved them once autosave assigned the real id, so the sidebar row showed no spinner and
Stop could not reach the generation, which kept holding a slot.
adoptDefaultThreadRun re-keys the run maps onto the real id from the thread adapter's
initialize(), where the id first exists; anything already filed under that id wins, since
that is a later run. The adapter captures its key once at run start, so it now resolves
the live key per use through runKeyForOwner, looking its own serverCancel up in the owner
map. Without that the migrated entries are stranded and the spinner never clears.
The denoising canvas was one global slot, so two diffusion chats overwrote each other and
the ownership tag then hid the visible preview until that thread emitted again. It is now
activeDiffusionCanvasByThreadId, written and cleared per thread, and the frame no longer
carries a threadId of its own. The bubble reads threadListItem.remoteId, dropping the dead
threadListItem.id arm: the writer tags unstable_threadId, which is exactly remoteId.
Two existing backend tests needed the same treatment. _bare_orchestrator skips __init__, so
it now sets the claim bookkeeping the worker ownership check reads. The Anthropic
passthrough gate test anchored on comment prose that a rewrap had broken; it anchors on the
code instead.
* Studio: hand the worker over cleanly between generations, and stop unresolved runs sharing each other's state
Worker ownership moved off the consumer and onto the dispatcher. Consumers read their
mailbox whenever they get around to it, so a request whose gen_done had been routed still
owned the worker while the next one ran, and a late Stop for it cancelled that one. The
dispatcher is the only place responses arrive in the order the worker produced them: it
now retires a request at its terminal response and promotes the next one, and answering a
request makes it the sole executor, since the subprocess runs one generation at a time.
reserve()'s immediate path ignored the unpark tickets that _grant_waiters_locked already
honours, so a request arriving between a slot freeing and an approved chat's next poll
took it, repeatedly. It applies the same reservation now.
Three places let concurrent first turns share state through the "__default" key. Nothing
links a run filed there to the id its thread later receives, so rather than guess, each
now declines when the key is ambiguous: adoption only re-keys a lone run, the composer
badge only claims a lone status, and the tool-output fallback only applies to a thread
that is still running. That leaves two concurrent first turns where they were before
adoption existed instead of handing one thread the other's handles.
A first turn's usage was never filed, because its key stayed null for the whole run while
autosave moved activeThreadId to the real id, so the context bar went blank after the
first reply. It resolves the adopted key like the cleanup handles do.
Cancelling a forced load left the UI with no model: the previous one stays resident until
/load's teardown, and the cancel path cleared the checkpoint without rolling back. It now
resyncs from the backend, which is right whether or not the load got that far.
The sidecar install drain is weighted 1:4 rather than halved, total unchanged. Only the
second half benefits from patience, and cutting it short refused installs whose chats had
already been stopped for nothing.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: give a first turn its real thread id before the run starts
A first turn filed every run handle under a shared unresolved key because
assistant-ui binds unstable_threadId before the thread is persisted. Two of them
overlapping there is unresolvable afterwards, and the last round's migration could
only decline rather than guess, which left neither sidebar row showing its run.
The id is available earlier than I claimed. append() already tracks
threadListItem.initialize() by the user message id, and createPersistedRunAdapter
already awaits that promise before invoking the adapter, so the thread is persisted
by the time the run begins. It was only being discarded: the tracked promise resolved
to void. It now resolves to the assigned id, and the wrapper hands it to the adapter
when assistant-ui had none. An id that is already set is never replaced, since that
would move a running chat's handles out from under the row watching them. The
existing unresolved-key guards stay as a safety net but should no longer carry weight.
The sidebar counted running thread ids rather than rows, so one compare conversation
read as two chats. It folds ids into rows through the same threadIds the row spinner
uses, and still counts a running id that matches no row.
_TrackedCancel always registered kind="chat", so an embeddings or raw completions
request appeared in the model-swap prompt as an unnamed conversation and confirming
cancelled it while calling it a chat. The non-conversation routes now pass their own
kind, and the prompt says "requests" whenever the snapshot is not all chats.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: withhold the shared worker cancel from a request the worker has left
Moving ownership to the dispatcher fixed reset_generation_state, but the token loop
signals the shared worker event directly and did not carry the same rule. A dispatched
consumer runs with mark_started off and can still be draining tokens buffered before
its gen_done was routed, so stopping it there ended whichever request the worker had
started next.
It now signals only when _owns_worker agrees, the same predicate reset_generation_state
uses. The local drain and return are unconditional, since those touch nothing but this
stream. The remaining _cancel_generation callers are deliberately global: subprocess
shutdown, the pre-load kill and unload_model.
* Studio: add the AGPL-3.0 header to the first-turn identity test
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: stop the dispatcher and a _gen_lock stream fighting over the response queue
Nothing stopped the dispatcher starting under a _gen_lock generation, so once compare
was opened while an ordinary chat was still streaming, both consumed _resp_queue and
whichever response the dispatcher took without a mailbox was dropped, gen_done included.
That chat truncated or hung. This PR is what makes it reachable, since navigating into
compare no longer ends the chat behind it.
Delaying the dispatcher would serialise compare behind whatever chat happens to be
streaming, so the direct readers get a mailbox instead. _direct_reader returns a reader,
a cancel drain and a release, and files the mailbox under _direct_mailboxes rather than
_mailboxes, which means "compare requests are in flight" to the unload and distributed
paths and must not count an ordinary chat.
Both directions close. The dispatcher finds the direct reader's mailbox instead of
dropping. And this reader can already be blocked on the queue when a compare request's
dispatcher starts, so a response that is not ours goes to its own mailbox rather than
being consumed, which would have corrupted the chat and hung the pane. All three
_gen_lock readers use it, and the cancel drain goes through it too.
The sidebar's return target still picked a raw pane id while the count grouped by row,
and /chat addresses compare with `compare`, not `thread`. It resolves through the same
items now, so a running compare row returns to its pair.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: keep worker ownership honest across audio, API traffic and a replaced worker
The audio-input send got a mailbox last round but stayed unclaimed, so a compare request
queued behind it looked like the oldest owner and stopping that queued request signalled
the shared event into the audio chat. It claims under the send lock and releases in the
finally, like _generate_inner.
Ownership is keyed on cancel-event identity with nothing tying it to a worker generation,
so a consumer still blocked on its mailbox when the process was replaced stayed recorded
as the executor, and a generation on the fresh worker could not be stopped.
_shutdown_subprocess clears that state once the process is confirmed dead, mailboxes
included: nothing routes to them again, and a stale one reads as compare activity to the
unload path. Not on the survived-SIGKILL path, which keeps its handle on purpose.
The four public /v1/messages trackers were registering as chats. The distinction is a
Studio thread, not the protocol, and those branches already say "No thread_id: public API
surface" while the Studio path passes payload.thread_id separately. They carry their own
kind now, so the swap prompt stops calling an external request a chat.
The swap confirmation still counted raw pane ids, so a compare conversation asked to stop
two chats and listed its title twice. It folds panes onto pairId and lowers the count by
what it collapsed, leaving a first turn the backend can count but not name.
Deep Research set runningByThreadId but registered no server-cancel handle, and that map
is how Stop, archive and delete reach a thread that is no longer active. Leaving the
outgoing thread running is this PR's doing, so the run was left unreachable while its
supervisor kept working against a conversation the user could delete.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: tighten the parallel-chats comments
* Studio: replay a Deep Research stop that arrived before the run existed
The handle is registered before createResearchRun resolves because the thread can be
stopped while that request is in flight, but it had no id to act on and dropped the stop.
The supervisor then followed a run the user had already stopped, archived or deleted.
It latches instead: a stop with no id yet sets a flag, and the adapter replays it against
the id the moment creation returns rather than starting to follow.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: fix worker ownership on a raced reroute, and the stop-chats prompt
Four review findings on the parallel-chats work, all reproduced first.
- _direct_reader hands a foreign response to its own mailbox, but skipped the
ownership move the dispatcher makes. A _gen_lock reader already blocked on
resp_queue can beat the compare dispatcher to that request's first response,
and the compare consumer opts out of marking, so nothing promoted it: the
direct request stayed the recorded executor, its late reset cancelled the
compare generation, and the compare chat's own Stop was ignored.
- A chat stopped while queued on _gen_lock was still claimed and sent once the
lock freed. Cancellation is only checked on a token, so a long prefill, or a
generation reaching gen_done without one, occupied the worker after Stop.
Same hole in the audio-input path, which shares the lock.
- The stop-chats prompt counted generation handles, not conversations. One chat
holds several while a tool continuation registers its next leg before the
previous unwinds, so it offered to stop two chats and listed one title.
- Ejecting a model confirms through that dialog, which told the user
"Unloading the model reloads the model" and offered "Stop and reload".
Confirming calls /unload and leaves nothing loaded.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: name the TTS run's thread so the stop prompt counts it once
The audio branch registers its run locally under the thread key but sent no
thread_id, so the backend tracker filed the same generation under no thread.
The stop-chats prompt then had a named local run and an unnamed backend one and,
since e8e7594 started adding unnamed entries to the named ones, counted a single
TTS chat as two requests. The backend already reads payload.thread_id, so
sending it lines both registries up on the same run.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2531 lines
103 KiB
Python
2531 lines
103 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Core inference backend."""
|
|
|
|
from unsloth import FastLanguageModel, FastVisionModel
|
|
from unsloth.chat_templates import get_chat_template
|
|
from transformers import TextIteratorStreamer, TextStreamer
|
|
from peft import PeftModel, PeftModelForCausalLM
|
|
|
|
import contextlib
|
|
import json
|
|
import sys
|
|
import torch
|
|
from pathlib import Path
|
|
from typing import Optional, Union, Generator, Tuple
|
|
from utils.models import ModelConfig, get_base_model_from_lora
|
|
from utils.paths import is_model_cached
|
|
from utils.transformers_dtype import dtype_kwargs
|
|
from utils.utils import format_error_message
|
|
from utils.hardware import (
|
|
get_device,
|
|
clear_gpu_cache,
|
|
log_gpu_memory,
|
|
get_device_map,
|
|
raise_if_offloaded,
|
|
get_visible_gpu_count,
|
|
)
|
|
from core.inference.audio_codecs import AudioCodecManager
|
|
from core.inference.runtime_context import runtime_context_length
|
|
from core.inference.message_content import content_to_text
|
|
from core.inference.chat_eos import (
|
|
chat_eos_repair,
|
|
resolve_chat_turn_end_eos_ids_using,
|
|
)
|
|
from core.inference.chat_template_helpers import (
|
|
ReasoningChannelNormalizer,
|
|
detect_reasoning_channel_markers,
|
|
detect_think_prefill,
|
|
)
|
|
from core.inference.presence_penalty import _make_presence_penalty_processor
|
|
from io import StringIO
|
|
import structlog
|
|
from loggers import get_logger
|
|
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class HarmonyTextStreamer:
|
|
"""Streaming text decoder for the gpt-oss harmony channel protocol.
|
|
|
|
gpt-oss emits multi-channel output via ``<|channel|>analysis<|message|>...``
|
|
/ ``<|channel|>final<|message|>...``. Plain skip_special_tokens streaming
|
|
glues channel names to content. This decodes with skip_special_tokens=False
|
|
and parses statefully: emit ``<think>`` on first analysis, stream analysis,
|
|
emit ``</think>`` on first final, stream final. Tracking per-channel lengths
|
|
avoids the delta-on-transformed bug where wrapping tags shift position.
|
|
Same put/end/iterator interface as TextIteratorStreamer.
|
|
"""
|
|
|
|
import re as _re
|
|
|
|
_HARMONY_RE = _re.compile(
|
|
r"<\|channel\|>(\w+)<\|message\|>(.*?)(?=<\|end\|>|<\|channel\|>|\Z)",
|
|
_re.DOTALL,
|
|
)
|
|
|
|
def __init__(
|
|
self,
|
|
tokenizer,
|
|
*,
|
|
skip_prompt: bool = True,
|
|
timeout: float = 0.2,
|
|
):
|
|
import queue
|
|
|
|
self.tokenizer = tokenizer
|
|
self.skip_prompt = skip_prompt
|
|
self.timeout = timeout
|
|
|
|
self._queue: queue.Queue = queue.Queue()
|
|
self._token_ids: list = []
|
|
self._prompt_len: int = 0
|
|
self._is_first_put: bool = True
|
|
self._stop: bool = False
|
|
|
|
# Stateful channel tracking avoids delta-on-transformed bugs
|
|
self._emitted_think_open: bool = False
|
|
self._emitted_think_close: bool = False
|
|
self._analysis_emitted: int = 0 # chars of analysis content emitted
|
|
self._final_emitted: int = 0 # chars of final content emitted
|
|
|
|
# put / end — called from the generation thread
|
|
|
|
def put(self, value):
|
|
"""Receive new token IDs from model.generate()."""
|
|
import torch
|
|
|
|
if isinstance(value, torch.Tensor):
|
|
# shape (batch, seq) — take first batch element
|
|
ids = value[0].tolist() if value.dim() > 1 else value.tolist()
|
|
elif isinstance(value, (list, tuple)):
|
|
ids = list(value)
|
|
else:
|
|
ids = [value]
|
|
|
|
if self._is_first_put and self.skip_prompt:
|
|
# First call is the full prompt; remember its length.
|
|
self._prompt_len = len(ids)
|
|
self._token_ids = list(ids)
|
|
self._is_first_put = False
|
|
return
|
|
|
|
self._token_ids.extend(ids)
|
|
|
|
# Decode only the generated part (after the prompt).
|
|
gen_ids = self._token_ids[self._prompt_len :]
|
|
raw = self.tokenizer.decode(gen_ids, skip_special_tokens = False)
|
|
self._process_incremental(raw)
|
|
|
|
def end(self):
|
|
"""Signal generation is complete."""
|
|
# Final decode to capture remaining content.
|
|
gen_ids = self._token_ids[self._prompt_len :]
|
|
if gen_ids:
|
|
raw = self.tokenizer.decode(gen_ids, skip_special_tokens = False)
|
|
self._process_incremental(raw)
|
|
|
|
# Close any open think tags.
|
|
if self._emitted_think_open and not self._emitted_think_close:
|
|
self._queue.put("</think>")
|
|
self._emitted_think_close = True
|
|
|
|
self._stop = True
|
|
self._queue.put(None) # sentinel
|
|
|
|
# Iterator interface — consumed by the streaming loop
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def __next__(self):
|
|
from queue import Empty
|
|
while True:
|
|
try:
|
|
val = self._queue.get(timeout = self.timeout)
|
|
except Empty:
|
|
if self._stop:
|
|
raise StopIteration
|
|
raise # propagate Empty so caller can check thread liveness
|
|
if val is None:
|
|
raise StopIteration
|
|
return val
|
|
|
|
# Stateful incremental harmony protocol parsing
|
|
|
|
def _process_incremental(self, raw: str) -> None:
|
|
"""Parse harmony channels and emit per-channel deltas (tracked by length, not whole-text diff)."""
|
|
# If raw has <|channel|> but no complete channel+message pair yet, buffer.
|
|
has_channel_token = "<|channel|>" in raw
|
|
matches = list(self._HARMONY_RE.finditer(raw))
|
|
|
|
if has_channel_token and not matches:
|
|
# Partial harmony markup still building — wait for more tokens.
|
|
return
|
|
|
|
if not has_channel_token and not matches:
|
|
return
|
|
|
|
for m in matches:
|
|
channel = m.group(1).lower()
|
|
content = m.group(2)
|
|
|
|
if channel == "analysis":
|
|
if not self._emitted_think_open:
|
|
self._queue.put("<think>")
|
|
self._emitted_think_open = True
|
|
|
|
new_content = content[self._analysis_emitted :]
|
|
if new_content:
|
|
self._analysis_emitted = len(content)
|
|
self._queue.put(new_content)
|
|
|
|
elif channel in ("final", "assistant"):
|
|
if self._emitted_think_open and not self._emitted_think_close:
|
|
self._queue.put("</think>")
|
|
self._emitted_think_close = True
|
|
|
|
new_content = content[self._final_emitted :]
|
|
if new_content:
|
|
self._final_emitted = len(content)
|
|
self._queue.put(new_content)
|
|
|
|
|
|
class ReasoningTextIteratorStreamer(TextIteratorStreamer):
|
|
"""TextIteratorStreamer that preserves native channel tokens until parsed."""
|
|
|
|
def __init__(
|
|
self,
|
|
tokenizer,
|
|
*,
|
|
markers: tuple[str, str],
|
|
skip_prompt: bool = True,
|
|
timeout: float = 0.2,
|
|
cancel_event = None,
|
|
**decode_kwargs,
|
|
):
|
|
decode_kwargs["skip_special_tokens"] = False
|
|
super().__init__(tokenizer, skip_prompt = skip_prompt, timeout = timeout, **decode_kwargs)
|
|
self._normalizer = ReasoningChannelNormalizer(*markers)
|
|
self._cancel_event = cancel_event
|
|
self._aborted = False
|
|
|
|
def abort(self):
|
|
"""Mark generation as failed so ``end`` drains without closing."""
|
|
self._aborted = True
|
|
|
|
def on_finalized_text(
|
|
self,
|
|
text: str,
|
|
stream_end: bool = False,
|
|
):
|
|
"""Queue canonical deltas, closing only on natural stream completion."""
|
|
delta = self._normalizer.feed(text)
|
|
if delta:
|
|
self.text_queue.put(delta, timeout = self.timeout)
|
|
|
|
if stream_end:
|
|
cancelled = self._aborted or (
|
|
self._cancel_event is not None and self._cancel_event.is_set()
|
|
)
|
|
tail = self._normalizer.drain() if cancelled else self._normalizer.finish()
|
|
if tail:
|
|
self.text_queue.put(tail, timeout = self.timeout)
|
|
self.text_queue.put(self.stop_signal, timeout = self.timeout)
|
|
|
|
|
|
class _GenerationThreadError(RuntimeError):
|
|
"""Generation worker failures that should propagate through stream routes."""
|
|
|
|
|
|
class InferenceBackend:
|
|
"""Unified inference backend supporting text, vision, and LoRA models"""
|
|
|
|
def __init__(self):
|
|
self.models = {}
|
|
self.active_model_name = None
|
|
self.loading_models = set()
|
|
self.loaded_local_models = [] # [(display_name, path), ...]
|
|
from core.inference.defaults import get_default_models
|
|
|
|
self.default_models = get_default_models()
|
|
self.device = get_device().value
|
|
self._audio_codec_manager = AudioCodecManager()
|
|
|
|
# _generation_lock serializes model.generate(). Plain Lock (NOT RLock):
|
|
# RLock reentrancy would let concurrent compare-mode requests race on
|
|
# the GPU. Acquired by the background generation thread, not the event-loop.
|
|
import threading
|
|
|
|
self._generation_lock = threading.Lock()
|
|
self._model_state_lock = threading.Lock()
|
|
|
|
logger.info(f"InferenceBackend initialized on {self.device}")
|
|
|
|
@staticmethod
|
|
def _normalize_top_k(top_k: int) -> int:
|
|
# API uses -1 to disable top-k; transformers uses 0.
|
|
return 0 if top_k < 0 else top_k
|
|
|
|
def _resolve_chat_eos(self, model_name: str) -> None:
|
|
"""Resolve this chat model's assistant-turn-end stop tokens once at load,
|
|
cache them in model_info, and repair generation_config so every
|
|
``.generate()`` path stops at the turn boundary.
|
|
|
|
Some checkpoints (e.g. Qwen3.5 / Qwen3.6 small chat models) end turns with
|
|
``<|im_end|>`` but ship ``config.eos_token_id = <|endoftext|>`` and no
|
|
``generation_config.json``, so paths that read ``generation_config`` (the
|
|
vision path, tool loops) run past the turn and loop. Turn-end markers are
|
|
derived from the chat_template (see chat_eos.resolve_chat_turn_end_eos_ids),
|
|
so base/coder models and harmony templates are left untouched.
|
|
"""
|
|
info = self.models.get(model_name) or {}
|
|
model = info.get("model")
|
|
container = info.get("tokenizer")
|
|
tokenizer = getattr(container, "tokenizer", container) # unwrap processors
|
|
if model is None or tokenizer is None:
|
|
return
|
|
# Vision models carry the chat_template on the processor, not the inner
|
|
# tokenizer. Read markers from whichever has one, but resolve ids on the
|
|
# generation tokenizer, else the vision path misses the turn-end token.
|
|
template_source = container if getattr(container, "chat_template", None) else tokenizer
|
|
try:
|
|
turn_end_ids = resolve_chat_turn_end_eos_ids_using(template_source, tokenizer)
|
|
except Exception as e: # never block a load on eos resolution
|
|
logger.warning("Chat turn-end eos resolution failed for %s: %s", model_name, e)
|
|
return
|
|
info["chat_turn_end_eos_ids"] = turn_end_ids
|
|
|
|
gen = getattr(model, "generation_config", None)
|
|
if gen is None:
|
|
return
|
|
repaired = chat_eos_repair(gen.eos_token_id, turn_end_ids)
|
|
if repaired is None:
|
|
return
|
|
previous = gen.eos_token_id
|
|
gen.eos_token_id = repaired
|
|
logger.info(
|
|
"Repaired generation_config.eos_token_id for %s: %s -> %s",
|
|
model_name,
|
|
previous,
|
|
repaired,
|
|
)
|
|
|
|
def load_model(
|
|
self,
|
|
config: ModelConfig,
|
|
max_seq_length: int = 2048,
|
|
dtype = None,
|
|
load_in_4bit: bool = True,
|
|
hf_token: Optional[str] = None,
|
|
trust_remote_code: bool = False,
|
|
gpu_ids: Optional[list[int]] = None,
|
|
) -> bool:
|
|
"""Load any model: base, LoRA adapter, text, or vision."""
|
|
# Keep the token so the native-template fallback can fetch a
|
|
# gated model's repo template later during generation.
|
|
self._hf_token = hf_token
|
|
# GGUF uses max_seq_length=0 as "model default"; Unsloth crashes on it.
|
|
if max_seq_length <= 0:
|
|
max_seq_length = 2048
|
|
|
|
try:
|
|
model_name = config.identifier
|
|
|
|
# Already loaded?
|
|
if model_name in self.models and self.models[model_name].get("model"):
|
|
logger.info(f"Model {model_name} already loaded")
|
|
if hf_token:
|
|
self.models[model_name]["hf_token"] = hf_token
|
|
self.active_model_name = model_name
|
|
return True
|
|
|
|
# Currently loading?
|
|
if model_name in self.loading_models:
|
|
logger.info(f"Model {model_name} is already being loaded")
|
|
return False
|
|
|
|
self.loading_models.add(model_name)
|
|
device_map = get_device_map(gpu_ids)
|
|
logger.info(
|
|
f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)"
|
|
)
|
|
|
|
self.models[model_name] = {
|
|
# Per-model token: the native-template fallback must use the
|
|
# token this model was loaded with, not whichever loaded last.
|
|
"hf_token": hf_token,
|
|
# Per-model consent: the native-template reload must re-use the
|
|
# exact trust_remote_code this model (and a LoRA's base) was loaded
|
|
# with, so a custom-code tokenizer repo can be re-fetched without
|
|
# executing any code the user did not already consent to.
|
|
"trust_remote_code": trust_remote_code,
|
|
"is_vision": config.is_vision,
|
|
"is_lora": config.is_lora,
|
|
"is_audio": config.is_audio,
|
|
"audio_type": config.audio_type,
|
|
"has_audio_input": config.has_audio_input,
|
|
"model_path": config.path,
|
|
"base_model": config.base_model if config.is_lora else None,
|
|
"loaded_adapters": {},
|
|
"active_adapter": None,
|
|
}
|
|
|
|
# ── Audio model loading path ──────────────────────────
|
|
if config.is_audio:
|
|
audio_type = config.audio_type
|
|
adapter_info = " (LoRA adapter)" if config.is_lora else ""
|
|
logger.info(f"Loading audio ({audio_type}) model{adapter_info}: {model_name}")
|
|
log_gpu_memory(f"Before loading {model_name}")
|
|
|
|
if audio_type == "csm":
|
|
from unsloth import FastModel
|
|
from transformers import CsmForConditionalGeneration
|
|
|
|
model, processor = FastModel.from_pretrained(
|
|
config.path,
|
|
auto_model = CsmForConditionalGeneration,
|
|
load_in_4bit = False,
|
|
device_map = device_map,
|
|
token = hf_token if hf_token and hf_token.strip() else None,
|
|
trust_remote_code = trust_remote_code,
|
|
)
|
|
FastModel.for_inference(model)
|
|
self.models[model_name]["model"] = model
|
|
self.models[model_name]["tokenizer"] = processor
|
|
self.models[model_name]["processor"] = processor
|
|
elif audio_type == "bicodec":
|
|
import os
|
|
from unsloth import FastModel
|
|
|
|
if config.is_lora and config.base_model:
|
|
# LoRA adapter: base_model is .../Spark-TTS-0.5B/LLM;
|
|
# BiCodec weights live in the parent dir.
|
|
base_path = config.base_model
|
|
if os.path.isdir(base_path):
|
|
abs_repo_path = os.path.abspath(os.path.dirname(base_path))
|
|
else:
|
|
# base_model is an HF ID — download it.
|
|
from huggingface_hub import snapshot_download
|
|
|
|
local_dir = base_path.split("/")[-1]
|
|
repo_path = snapshot_download(base_path, local_dir = local_dir)
|
|
abs_repo_path = os.path.abspath(repo_path)
|
|
|
|
logger.info(
|
|
f"Spark-TTS LoRA: loading adapter from {config.path}, BiCodec from {abs_repo_path}"
|
|
)
|
|
model, tokenizer = FastModel.from_pretrained(
|
|
config.path,
|
|
dtype = torch.float32,
|
|
load_in_4bit = False,
|
|
device_map = device_map,
|
|
token = hf_token if hf_token and hf_token.strip() else None,
|
|
trust_remote_code = trust_remote_code,
|
|
)
|
|
else:
|
|
# Base model: download full HF repo, load from /LLM subfolder
|
|
from huggingface_hub import snapshot_download
|
|
|
|
hf_repo = config.path
|
|
local_dir = hf_repo.split("/")[-1]
|
|
repo_path = snapshot_download(hf_repo, local_dir = local_dir)
|
|
abs_repo_path = os.path.abspath(repo_path)
|
|
llm_path = os.path.join(abs_repo_path, "LLM")
|
|
logger.info(
|
|
f"Spark-TTS: downloaded repo to {repo_path}, loading LLM from {llm_path}"
|
|
)
|
|
|
|
model, tokenizer = FastModel.from_pretrained(
|
|
llm_path,
|
|
dtype = torch.float32,
|
|
load_in_4bit = False,
|
|
device_map = device_map,
|
|
token = hf_token if hf_token and hf_token.strip() else None,
|
|
trust_remote_code = trust_remote_code,
|
|
)
|
|
|
|
FastModel.for_inference(model)
|
|
self.models[model_name]["model"] = model
|
|
self.models[model_name]["tokenizer"] = tokenizer
|
|
self.models[model_name]["model_repo_path"] = abs_repo_path
|
|
elif audio_type == "dac":
|
|
# OuteTTS uses FastModel (not FastLanguageModel)
|
|
from unsloth import FastModel
|
|
|
|
model, tokenizer = FastModel.from_pretrained(
|
|
config.path,
|
|
max_seq_length = max_seq_length,
|
|
load_in_4bit = False,
|
|
device_map = device_map,
|
|
token = hf_token if hf_token and hf_token.strip() else None,
|
|
trust_remote_code = trust_remote_code,
|
|
)
|
|
FastModel.for_inference(model)
|
|
self.models[model_name]["model"] = model
|
|
self.models[model_name]["tokenizer"] = tokenizer
|
|
elif audio_type == "whisper":
|
|
# Whisper ASR — uses FastModel with WhisperForConditionalGeneration
|
|
from unsloth import FastModel
|
|
from transformers import WhisperForConditionalGeneration
|
|
|
|
model, tokenizer = FastModel.from_pretrained(
|
|
config.path,
|
|
auto_model = WhisperForConditionalGeneration,
|
|
whisper_language = "English",
|
|
whisper_task = "transcribe",
|
|
load_in_4bit = False,
|
|
device_map = device_map,
|
|
token = hf_token if hf_token and hf_token.strip() else None,
|
|
trust_remote_code = trust_remote_code,
|
|
)
|
|
FastModel.for_inference(model)
|
|
model.eval()
|
|
|
|
# ASR pipeline (per notebook)
|
|
from transformers import pipeline as hf_pipeline
|
|
|
|
whisper_pipe = hf_pipeline(
|
|
"automatic-speech-recognition",
|
|
model = model,
|
|
tokenizer = tokenizer.tokenizer,
|
|
feature_extractor = tokenizer.feature_extractor,
|
|
processor = tokenizer,
|
|
return_language = True,
|
|
**dtype_kwargs(torch.float16),
|
|
)
|
|
self.models[model_name]["model"] = model
|
|
self.models[model_name]["tokenizer"] = tokenizer
|
|
self.models[model_name]["whisper_pipeline"] = whisper_pipe
|
|
else:
|
|
# SNAC (Orpheus) uses FastLanguageModel
|
|
model, tokenizer = FastLanguageModel.from_pretrained(
|
|
model_name = config.path,
|
|
max_seq_length = max_seq_length,
|
|
load_in_4bit = False,
|
|
device_map = device_map,
|
|
token = hf_token if hf_token and hf_token.strip() else None,
|
|
trust_remote_code = trust_remote_code,
|
|
)
|
|
FastLanguageModel.for_inference(model)
|
|
self.models[model_name]["model"] = model
|
|
self.models[model_name]["tokenizer"] = tokenizer
|
|
|
|
# Load external codec for TTS audio types
|
|
# (Whisper is ASR, audio_vlm is audio input — neither needs one)
|
|
if audio_type not in ("whisper", "audio_vlm"):
|
|
model_repo_path = self.models[model_name].get("model_repo_path")
|
|
self._audio_codec_manager.load_codec(
|
|
audio_type, self.device, model_repo_path = model_repo_path
|
|
)
|
|
|
|
# Reject CPU/disk offload for audio models too
|
|
raise_if_offloaded(self.models[model_name]["model"], device_map, "Inference")
|
|
self.models[model_name]["context_length"] = runtime_context_length(
|
|
self.models[model_name].get("model"),
|
|
max_seq_length,
|
|
)
|
|
|
|
self.active_model_name = model_name
|
|
self.loading_models.discard(model_name)
|
|
logger.info(f"Successfully loaded audio model: {model_name}")
|
|
log_gpu_memory(f"After loading {model_name}")
|
|
return True
|
|
|
|
model_type = "vision" if config.is_vision else "text"
|
|
adapter_info = " (LoRA adapter)" if self.models[model_name]["is_lora"] else ""
|
|
logger.info(f"Loading {model_type} model{adapter_info}: {model_name}")
|
|
log_gpu_memory(f"Before loading {model_name}")
|
|
|
|
# Same load path for base models and LoRA adapters
|
|
if config.is_vision:
|
|
# Vision model (or vision LoRA adapter)
|
|
model, processor = FastVisionModel.from_pretrained(
|
|
model_name = config.path, # Can be base model OR LoRA adapter path
|
|
max_seq_length = max_seq_length,
|
|
dtype = dtype,
|
|
load_in_4bit = load_in_4bit,
|
|
device_map = device_map,
|
|
token = hf_token if hf_token and hf_token.strip() else None,
|
|
trust_remote_code = trust_remote_code,
|
|
)
|
|
|
|
FastVisionModel.for_inference(model)
|
|
|
|
# FastVisionModel may return a raw tokenizer instead of a
|
|
# Processor for some models (e.g. Gemma-3); load the real one.
|
|
from transformers import ProcessorMixin
|
|
|
|
if not (
|
|
isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor")
|
|
):
|
|
# LoRA adapters: use base model. Local merged exports: read base from export_metadata.json.
|
|
processor_source = config.base_model if config.is_lora else config.identifier
|
|
if not config.is_lora and config.is_local:
|
|
_meta_path = Path(config.path) / "export_metadata.json"
|
|
try:
|
|
if _meta_path.exists():
|
|
_meta = json.loads(_meta_path.read_text(encoding = "utf-8"))
|
|
if _meta.get("base_model"):
|
|
processor_source = _meta["base_model"]
|
|
except Exception:
|
|
pass
|
|
logger.warning(
|
|
f"FastVisionModel returned {type(processor).__name__} (no image_processor) "
|
|
f"for '{model_name}' — loading proper processor from '{processor_source}'"
|
|
)
|
|
from transformers import AutoProcessor
|
|
|
|
processor = AutoProcessor.from_pretrained(
|
|
processor_source,
|
|
token = hf_token if hf_token and hf_token.strip() else None,
|
|
trust_remote_code = trust_remote_code,
|
|
)
|
|
logger.info(f"Loaded {type(processor).__name__} from {processor_source}")
|
|
|
|
self.models[model_name]["model"] = model
|
|
self.models[model_name]["tokenizer"] = processor
|
|
self.models[model_name]["processor"] = processor
|
|
|
|
else:
|
|
# Text model (or text LoRA adapter)
|
|
model, tokenizer = FastLanguageModel.from_pretrained(
|
|
model_name = config.path, # Can be base model OR LoRA adapter path
|
|
max_seq_length = max_seq_length,
|
|
dtype = dtype,
|
|
load_in_4bit = load_in_4bit,
|
|
device_map = device_map,
|
|
token = hf_token if hf_token and hf_token.strip() else None,
|
|
trust_remote_code = trust_remote_code,
|
|
)
|
|
|
|
FastLanguageModel.for_inference(model)
|
|
|
|
self.models[model_name]["model"] = model
|
|
self.models[model_name]["tokenizer"] = tokenizer
|
|
|
|
raise_if_offloaded(self.models[model_name]["model"], device_map, "Inference")
|
|
self.models[model_name]["context_length"] = runtime_context_length(
|
|
self.models[model_name].get("model"),
|
|
max_seq_length,
|
|
)
|
|
|
|
self._resolve_chat_eos(model_name)
|
|
self._load_chat_template_info(model_name)
|
|
|
|
self.active_model_name = model_name
|
|
self.loading_models.discard(model_name)
|
|
|
|
logger.info(f"Successfully loaded model: {model_name}")
|
|
log_gpu_memory(f"After loading {model_name}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to load model: {e}")
|
|
error_msg = format_error_message(e, config.identifier)
|
|
|
|
# Cleanup on failure
|
|
if model_name in self.models:
|
|
del self.models[model_name]
|
|
self.loading_models.discard(model_name)
|
|
|
|
raise Exception(error_msg)
|
|
|
|
def unload_model(self, model_name: str) -> bool:
|
|
"""Remove a model from the registry and clear GPU memory."""
|
|
if model_name in self.models:
|
|
try:
|
|
# Clean up codecs for audio models
|
|
if self.models[model_name].get("is_audio"):
|
|
self._audio_codec_manager.unload()
|
|
|
|
logger.info(f"Unloading model '{model_name}' from memory.")
|
|
del self.models[model_name]
|
|
|
|
# Clear the active model if it was the one unloaded
|
|
if self.active_model_name == model_name:
|
|
self.active_model_name = None
|
|
|
|
clear_gpu_cache()
|
|
|
|
# Drop stale compiled cache for the next model. On spawn platforms,
|
|
# preserve trainer files so concurrent dataset.map() workers can import them.
|
|
import sys as _sys
|
|
from utils.cache_cleanup import clear_unsloth_compiled_cache
|
|
|
|
_preserve = ["Unsloth*Trainer.py"] if _sys.platform in ("win32", "darwin") else None
|
|
clear_unsloth_compiled_cache(preserve_patterns = _preserve)
|
|
|
|
logger.info(f"Model '{model_name}' successfully unloaded.")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Error while unloading model '{model_name}': {e}")
|
|
return False
|
|
else:
|
|
logger.warning(
|
|
f"Attempted to unload model '{model_name}', but it was not found in the registry."
|
|
)
|
|
return True
|
|
|
|
def revert_to_base_model(self, base_model_name: str) -> bool:
|
|
"""Revert the model to its pristine base state by unloading and
|
|
deleting all adapter configurations."""
|
|
if base_model_name not in self.models:
|
|
return False
|
|
|
|
model = self.models[base_model_name].get("model")
|
|
|
|
try:
|
|
# Unload adapter weights if model is a PeftModel.
|
|
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
|
|
logger.info(f"Unloading LoRA adapters from '{base_model_name}'...")
|
|
unwrapped_base_model = model.unload()
|
|
self.models[base_model_name]["model"] = unwrapped_base_model
|
|
model = unwrapped_base_model
|
|
|
|
# model.unload() can leave a peft_config; removing it avoids
|
|
# "multiple adapters" warnings on the next from_pretrained().
|
|
if hasattr(model, "peft_config"):
|
|
del model.peft_config
|
|
|
|
logger.info(f"Model '{base_model_name}' reverted to clean base state.")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to revert model to base state: {e}")
|
|
import traceback
|
|
|
|
logger.error(traceback.format_exc())
|
|
return False
|
|
|
|
def load_for_eval(
|
|
self,
|
|
lora_path: str,
|
|
max_seq_length: int = 2048,
|
|
dtype = None,
|
|
load_in_4bit: bool = True,
|
|
hf_token: Optional[str] = None,
|
|
gpu_ids: Optional[list[int]] = None,
|
|
) -> Tuple[bool, Optional[str], Optional[str]]:
|
|
"""Ensure the base model and the given adapter are loaded.
|
|
Idempotent and handles all states correctly.
|
|
"""
|
|
try:
|
|
from utils.models import ModelConfig
|
|
|
|
lora_config = ModelConfig.from_lora_path(lora_path, hf_token)
|
|
if not lora_config:
|
|
return False, None, None
|
|
|
|
base_model_name = lora_config.base_model
|
|
|
|
# 1. Load the base model if not already in memory
|
|
if base_model_name not in self.models or not self.models[base_model_name].get("model"):
|
|
logger.info(f"Base model '{base_model_name}' not loaded, loading now.")
|
|
base_config = ModelConfig.from_ui_selection(base_model_name, None, is_lora = False)
|
|
if not self.load_model(
|
|
base_config,
|
|
max_seq_length,
|
|
dtype,
|
|
load_in_4bit,
|
|
hf_token,
|
|
gpu_ids = gpu_ids,
|
|
):
|
|
return False, None, None
|
|
|
|
self.active_model_name = base_model_name
|
|
|
|
# 2. Derive adapter name from the user's selection
|
|
adapter_name = lora_path.split("/")[-1].replace(".", "_")
|
|
|
|
# 3. Ensure this adapter is loaded (load_adapter only reads from
|
|
# disk if the model doesn't already have it).
|
|
adapter_success = self.load_adapter(
|
|
base_model_name = base_model_name,
|
|
adapter_path = lora_path,
|
|
adapter_name = adapter_name,
|
|
)
|
|
if not adapter_success:
|
|
return False, base_model_name, None
|
|
|
|
# 4. Return the verified adapter name for the UI.
|
|
return True, base_model_name, adapter_name
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error during load_for_eval: {e}")
|
|
import traceback
|
|
|
|
logger.error(traceback.format_exc())
|
|
return False, None, None
|
|
|
|
def load_adapter(self, base_model_name: str, adapter_path: str, adapter_name: str) -> bool:
|
|
"""Load an adapter onto the model only if not already attached."""
|
|
model = self.models[base_model_name].get("model")
|
|
|
|
# Most reliable check: adapter name already in the model's config.
|
|
if hasattr(model, "peft_config") and adapter_name in model.peft_config:
|
|
logger.info(
|
|
f"Adapter '{adapter_name}' is already attached to the model. Skipping load."
|
|
)
|
|
return True
|
|
|
|
try:
|
|
logger.info(
|
|
f"Loading new adapter '{adapter_name}' from '{adapter_path}' onto {base_model_name}"
|
|
)
|
|
model.load_adapter(adapter_path, adapter_name = adapter_name)
|
|
|
|
# Update the registry only after a successful load.
|
|
if "loaded_adapters" not in self.models[base_model_name]:
|
|
self.models[base_model_name]["loaded_adapters"] = {}
|
|
self.models[base_model_name]["loaded_adapters"][adapter_name] = adapter_path
|
|
|
|
total_adapters = len(getattr(model, "peft_config", {}))
|
|
logger.info(
|
|
f"Adapter '{adapter_name}' loaded successfully. (Total unique adapters on model: {total_adapters})"
|
|
)
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to load adapter '{adapter_name}': {e}")
|
|
return False
|
|
|
|
def set_active_adapter(self, base_model_name: str, adapter_name: str) -> bool:
|
|
"""Set the active adapter for generation."""
|
|
model = self.models[base_model_name].get("model")
|
|
try:
|
|
logger.info(f"Setting active adapter to: '{adapter_name}'")
|
|
model.set_adapter(adapter_name)
|
|
self.models[base_model_name]["active_adapter"] = adapter_name
|
|
return True
|
|
except Exception as e:
|
|
# Catches "adapter not found" if something goes wrong.
|
|
logger.error(f"Failed to set active adapter to '{adapter_name}': {e}")
|
|
return False
|
|
|
|
def _apply_adapter_state(self, use_adapter: Optional[Union[bool, str]]) -> None:
|
|
"""Apply adapter state before generation (must hold _generation_lock).
|
|
|
|
Toggles PEFT enable/disable_adapter_layers (non-destructive, no reload).
|
|
use_adapter: None = no change, False = base model, True = current adapter,
|
|
str = named adapter.
|
|
"""
|
|
if use_adapter is None:
|
|
return
|
|
|
|
base = self.active_model_name
|
|
if not base or base not in self.models:
|
|
return
|
|
|
|
model_info = self.models[base]
|
|
model = model_info.get("model")
|
|
if model is None:
|
|
return
|
|
|
|
if use_adapter is False:
|
|
# Disable LoRA layers -> base model output.
|
|
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
|
|
logger.info(
|
|
f"Compare mode: disabling adapters on '{base}' for base model generation"
|
|
)
|
|
model.base_model.disable_adapter_layers()
|
|
else:
|
|
logger.info(f"Compare mode: model '{base}' is not a PeftModel, already base")
|
|
|
|
elif use_adapter is True:
|
|
# Re-enable LoRA layers -> adapter output.
|
|
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
|
|
logger.info(f"Compare mode: enabling adapters on '{base}' for LoRA generation")
|
|
model.base_model.enable_adapter_layers()
|
|
else:
|
|
logger.warning("use_adapter=true but model is not a PeftModel")
|
|
|
|
elif isinstance(use_adapter, str):
|
|
# Enable adapters and set the named one active.
|
|
if isinstance(model, (PeftModel, PeftModelForCausalLM)):
|
|
logger.info(f"Compare mode: enabling adapter '{use_adapter}' on '{base}'")
|
|
model.base_model.enable_adapter_layers()
|
|
self.set_active_adapter(base, use_adapter)
|
|
else:
|
|
logger.warning(f"use_adapter='{use_adapter}' but model is not a PeftModel")
|
|
|
|
def generate_with_adapter_control(
|
|
self,
|
|
use_adapter: Optional[Union[bool, str]] = None,
|
|
cancel_event = None,
|
|
**gen_kwargs,
|
|
) -> Generator[str, None, None]:
|
|
"""Thread-safe generation with optional adapter toggling.
|
|
|
|
Adapter toggle + model.generate() are serialized by _generation_lock in
|
|
the background thread, avoiding the RLock-reentrant race when two async
|
|
SSE handlers share one event-loop thread. use_adapter: see _apply_adapter_state.
|
|
"""
|
|
yield from self._generate_chat_response_inner(
|
|
cancel_event = cancel_event, _adapter_state = use_adapter, **gen_kwargs
|
|
)
|
|
|
|
def generate_chat_completion_with_tools(
|
|
self,
|
|
messages: list,
|
|
tools: list,
|
|
system_prompt: str = "",
|
|
temperature: float = 0.7,
|
|
top_p: float = 0.9,
|
|
top_k: int = 40,
|
|
min_p: float = 0.0,
|
|
max_new_tokens: int = 2048,
|
|
repetition_penalty: float = 1.0,
|
|
cancel_event = None,
|
|
enable_thinking: Optional[bool] = None,
|
|
reasoning_effort: Optional[str] = None,
|
|
preserve_thinking: Optional[bool] = None,
|
|
max_tool_iterations: int = 25,
|
|
auto_heal_tool_calls: bool = True,
|
|
nudge_tool_calls: Optional[bool] = None,
|
|
tool_call_timeout: int = 300,
|
|
session_id: Optional[str] = None,
|
|
thread_id: Optional[str] = None,
|
|
rag_scope: Optional[dict] = None,
|
|
presence_penalty: float = 0.0,
|
|
reasoning_prefilled: bool = False,
|
|
):
|
|
"""Run an agentic tool loop on top of ``generate_chat_response``.
|
|
|
|
Yields the same event-dict protocol as the GGUF path so the route
|
|
layer can stream both backends through one helper. Each event is one of:
|
|
|
|
* ``{"type": "status", "text": ...}``
|
|
* ``{"type": "content", "text": cumulative_text}``
|
|
* ``{"type": "tool_start", "tool_name", "tool_call_id", "arguments"}``
|
|
* ``{"type": "tool_end", "tool_name", "tool_call_id", "result"}``
|
|
"""
|
|
from core.inference.safetensors_agentic import run_safetensors_tool_loop
|
|
from core.inference.tools import execute_tool
|
|
|
|
def _single_turn(conv: list, *, active_tools: Optional[list[dict]] = None):
|
|
# conv already has the system message -- avoid double-prepend.
|
|
# `active_tools` is supplied by run_safetensors_tool_loop so one-shot
|
|
# tools such as render_html can be removed from later same-response prompts.
|
|
turn_tools = active_tools if active_tools is not None else tools
|
|
yield from self._generate_chat_response_inner(
|
|
messages = conv,
|
|
system_prompt = "",
|
|
temperature = temperature,
|
|
top_p = top_p,
|
|
top_k = top_k,
|
|
min_p = min_p,
|
|
max_new_tokens = max_new_tokens,
|
|
repetition_penalty = repetition_penalty,
|
|
cancel_event = cancel_event,
|
|
tools = turn_tools,
|
|
enable_thinking = enable_thinking,
|
|
reasoning_effort = reasoning_effort,
|
|
preserve_thinking = preserve_thinking,
|
|
presence_penalty = presence_penalty,
|
|
)
|
|
|
|
initial = list(messages)
|
|
if system_prompt:
|
|
initial = [{"role": "system", "content": system_prompt}] + initial
|
|
|
|
yield from run_safetensors_tool_loop(
|
|
single_turn = _single_turn,
|
|
messages = initial,
|
|
tools = tools,
|
|
execute_tool = execute_tool,
|
|
cancel_event = cancel_event,
|
|
auto_heal_tool_calls = auto_heal_tool_calls,
|
|
nudge_tool_calls = nudge_tool_calls,
|
|
max_tool_iterations = max_tool_iterations,
|
|
tool_call_timeout = tool_call_timeout,
|
|
session_id = session_id,
|
|
thread_id = thread_id,
|
|
rag_scope = rag_scope,
|
|
reasoning_prefilled = reasoning_prefilled,
|
|
)
|
|
|
|
def generate_chat_response(
|
|
self,
|
|
messages: list,
|
|
system_prompt: str,
|
|
image = None,
|
|
temperature: float = 0.7,
|
|
top_p: float = 0.9,
|
|
top_k: int = 40,
|
|
min_p: float = 0.0,
|
|
max_new_tokens: int = 256,
|
|
repetition_penalty: float = 1.0,
|
|
cancel_event = None,
|
|
tools: Optional[list] = None,
|
|
enable_thinking: Optional[bool] = None,
|
|
reasoning_effort: Optional[str] = None,
|
|
preserve_thinking: Optional[bool] = None,
|
|
presence_penalty: float = 0.0,
|
|
) -> Generator[str, None, None]:
|
|
"""Generate response for text or vision models (lock held by background thread).
|
|
|
|
``tools`` / ``enable_thinking`` / ``reasoning_effort`` / ``preserve_thinking``
|
|
are forwarded into ``apply_chat_template`` so templates that understand them
|
|
(Qwen3, Llama 3.1+, gpt-oss harmony) advertise tool schemas / reasoning controls.
|
|
``presence_penalty`` matches the GGUF sampling path (0 disables it).
|
|
"""
|
|
yield from self._generate_chat_response_inner(
|
|
messages = messages,
|
|
system_prompt = system_prompt,
|
|
image = image,
|
|
temperature = temperature,
|
|
top_p = top_p,
|
|
top_k = top_k,
|
|
min_p = min_p,
|
|
max_new_tokens = max_new_tokens,
|
|
repetition_penalty = repetition_penalty,
|
|
cancel_event = cancel_event,
|
|
tools = tools,
|
|
enable_thinking = enable_thinking,
|
|
reasoning_effort = reasoning_effort,
|
|
preserve_thinking = preserve_thinking,
|
|
presence_penalty = presence_penalty,
|
|
)
|
|
|
|
def _generate_chat_response_inner(
|
|
self,
|
|
messages: list,
|
|
system_prompt: str = "",
|
|
image = None,
|
|
temperature: float = 0.7,
|
|
top_p: float = 0.9,
|
|
top_k: int = 40,
|
|
min_p: float = 0.0,
|
|
max_new_tokens: int = 256,
|
|
repetition_penalty: float = 1.0,
|
|
cancel_event = None,
|
|
_adapter_state = None,
|
|
tools: Optional[list] = None,
|
|
enable_thinking: Optional[bool] = None,
|
|
reasoning_effort: Optional[str] = None,
|
|
preserve_thinking: Optional[bool] = None,
|
|
presence_penalty: float = 0.0,
|
|
) -> Generator[str, None, None]:
|
|
"""Inner generation logic, called by generate_chat_response and
|
|
generate_with_adapter_control.
|
|
|
|
_adapter_state is passed to generate_stream/vision so the background
|
|
thread can toggle adapters under the generation lock.
|
|
"""
|
|
if not self.active_model_name:
|
|
raise RuntimeError("No active model")
|
|
|
|
model_info = self.models[self.active_model_name]
|
|
is_vision = model_info.get("is_vision", False)
|
|
tokenizer = model_info.get("tokenizer") or model_info.get("processor")
|
|
# Unwrap processor -> raw tokenizer for VLMs on the text path.
|
|
tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
|
|
top_k = self._normalize_top_k(top_k)
|
|
|
|
if is_vision and image:
|
|
# Verify the stored processor can handle images; FastVisionModel may
|
|
# return a raw tokenizer instead of a ProcessorMixin (e.g. Gemma-3).
|
|
from transformers import ProcessorMixin
|
|
|
|
processor = model_info.get("processor")
|
|
has_image_processing = processor is not None and (
|
|
isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor")
|
|
)
|
|
if has_image_processing:
|
|
yield from self._generate_vision_response(
|
|
messages,
|
|
system_prompt,
|
|
image,
|
|
temperature,
|
|
top_p,
|
|
top_k,
|
|
min_p,
|
|
max_new_tokens,
|
|
repetition_penalty,
|
|
cancel_event = cancel_event,
|
|
presence_penalty = presence_penalty,
|
|
)
|
|
return
|
|
else:
|
|
logger.warning(
|
|
f"Model '{self.active_model_name}' is marked as vision but its processor "
|
|
f"({type(processor).__name__}) has no image_processor — "
|
|
f"falling back to text-only generation (image will be ignored)."
|
|
)
|
|
|
|
# Text path: messages are already in ChatML format from eval.py.
|
|
|
|
# Step 1: apply get_chat_template if model is in mapper.
|
|
try:
|
|
from utils.datasets import (
|
|
MODEL_TO_TEMPLATE_MAPPER,
|
|
get_tokenizer_chat_template,
|
|
)
|
|
model_name_lower = self.active_model_name.lower()
|
|
|
|
if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
|
|
template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
|
|
logger.info(
|
|
f"Applying chat template '{template_name}' for {self.active_model_name}"
|
|
)
|
|
|
|
tokenizer = get_chat_template(
|
|
tokenizer,
|
|
chat_template = template_name,
|
|
)
|
|
# The mapper installs the effective template only now, at generate
|
|
# time, so re-resolve and UNION into the load-time cache (never
|
|
# overwrite). get_chat_template can return a remapped tokenizer
|
|
# (turn-end folded onto doc-eos) while generate_stream reads the
|
|
# original, so take marker strings from the mapped template but
|
|
# resolve their ids on the original.
|
|
try:
|
|
_gen_tok = model_info.get("tokenizer") or tokenizer
|
|
refreshed = resolve_chat_turn_end_eos_ids_using(
|
|
getattr(tokenizer, "tokenizer", tokenizer),
|
|
getattr(_gen_tok, "tokenizer", _gen_tok),
|
|
)
|
|
existing = model_info.get("chat_turn_end_eos_ids") or []
|
|
model_info["chat_turn_end_eos_ids"] = sorted(set(existing) | set(refreshed))
|
|
except Exception as e:
|
|
logger.warning(f"Could not refresh chat turn-end eos after template: {e}")
|
|
else:
|
|
logger.info(
|
|
f"No registered Unsloth template for {self.active_model_name}, using tokenizer default"
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"Could not apply get_chat_template: {e}")
|
|
|
|
# Step 2: format with tokenizer.apply_chat_template().
|
|
if system_prompt:
|
|
template_messages = [{"role": "system", "content": system_prompt}] + messages
|
|
else:
|
|
template_messages = messages
|
|
reasoning_channel_markers_resolved = False
|
|
try:
|
|
if not (hasattr(tokenizer, "chat_template") and tokenizer.chat_template):
|
|
raise ValueError(
|
|
f"Model '{self.active_model_name}' has no chat_template set in its "
|
|
f"tokenizer_config.json. This is usually a problem with the model's "
|
|
f"HuggingFace repository — it is missing a 'chat_template' key. "
|
|
f"Please use a model that includes a chat template, or manually set "
|
|
f"one via tokenizer.chat_template before inference."
|
|
)
|
|
reasoning_channel_markers = None
|
|
formatted_prompt = self._apply_chat_template_for_generation(
|
|
tokenizer,
|
|
template_messages,
|
|
tools = tools,
|
|
enable_thinking = enable_thinking,
|
|
reasoning_effort = reasoning_effort,
|
|
preserve_thinking = preserve_thinking,
|
|
)
|
|
|
|
# If tools were requested but the (possibly overridden) template ignored
|
|
# them, fall back to the model's native template (shared with MLX).
|
|
from core.inference.chat_template_helpers import (
|
|
render_with_native_template_fallback,
|
|
)
|
|
|
|
render_result = render_with_native_template_fallback(
|
|
formatted_prompt = formatted_prompt,
|
|
tokenizer = tokenizer,
|
|
model_info = model_info,
|
|
active_model_name = self.active_model_name,
|
|
messages = template_messages,
|
|
tools = tools,
|
|
enable_thinking = enable_thinking,
|
|
reasoning_effort = reasoning_effort,
|
|
preserve_thinking = preserve_thinking,
|
|
apply_fn = self._apply_chat_template_for_generation,
|
|
hf_token = model_info.get("hf_token"),
|
|
return_metadata = True,
|
|
)
|
|
formatted_prompt = render_result.prompt
|
|
reasoning_channel_markers = render_result.reasoning_channel_markers
|
|
reasoning_channel_markers_resolved = True
|
|
|
|
logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...")
|
|
except Exception as e:
|
|
logger.error(f"Error applying chat template: {e}")
|
|
# Fall back to manual formatting
|
|
formatted_prompt = self.format_chat_prompt(messages, system_prompt)
|
|
reasoning_channel_markers = None
|
|
reasoning_channel_markers_resolved = True
|
|
|
|
# Step 3: generate
|
|
yield from self.generate_stream(
|
|
formatted_prompt,
|
|
temperature,
|
|
top_p,
|
|
top_k,
|
|
min_p,
|
|
max_new_tokens,
|
|
repetition_penalty,
|
|
cancel_event = cancel_event,
|
|
_adapter_state = _adapter_state,
|
|
presence_penalty = presence_penalty,
|
|
reasoning_channel_markers = reasoning_channel_markers,
|
|
reasoning_channel_markers_resolved = reasoning_channel_markers_resolved,
|
|
)
|
|
|
|
def _generate_vision_response(
|
|
self,
|
|
messages,
|
|
system_prompt,
|
|
image,
|
|
temperature,
|
|
top_p,
|
|
top_k,
|
|
min_p,
|
|
max_new_tokens,
|
|
repetition_penalty,
|
|
cancel_event = None,
|
|
presence_penalty: float = 0.0,
|
|
) -> Generator[str, None, None]:
|
|
"""Handle vision model generation with true token-by-token streaming."""
|
|
model_info = self.models[self.active_model_name]
|
|
model = model_info["model"]
|
|
processor = model_info["processor"]
|
|
# FastVisionModel may return a raw tokenizer (e.g. GemmaTokenizerFast)
|
|
# for some models. Safe unwrap for tokenize-only ops.
|
|
raw_tokenizer = getattr(processor, "tokenizer", processor)
|
|
|
|
# Extract user message
|
|
user_message = ""
|
|
if messages and messages[-1]["role"] == "user":
|
|
import re
|
|
user_message = content_to_text(messages[-1]["content"])
|
|
user_message = re.sub(r"<img[^>]*>", "", user_message).strip()
|
|
|
|
if not user_message:
|
|
user_message = "Describe this image." if image else "Hello"
|
|
|
|
# Prepare vision messages
|
|
if image:
|
|
user_msg = {
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "image"},
|
|
{"type": "text", "text": user_message},
|
|
],
|
|
}
|
|
if system_prompt:
|
|
vision_messages = [
|
|
{
|
|
"role": "system",
|
|
"content": [{"type": "text", "text": system_prompt}],
|
|
},
|
|
user_msg,
|
|
]
|
|
else:
|
|
vision_messages = [user_msg]
|
|
|
|
try:
|
|
input_text = processor.apply_chat_template(
|
|
vision_messages, add_generation_prompt = True, tokenize = False
|
|
)
|
|
except Exception as e:
|
|
if system_prompt:
|
|
logger.warning(
|
|
f"Vision processor for '{self.active_model_name}' may not support "
|
|
f"system messages; retrying without. Original error: {e}"
|
|
)
|
|
vision_messages = [user_msg]
|
|
input_text = processor.apply_chat_template(
|
|
vision_messages, add_generation_prompt = True, tokenize = False
|
|
)
|
|
else:
|
|
raise
|
|
inputs = processor(
|
|
image,
|
|
input_text,
|
|
add_special_tokens = False,
|
|
return_tensors = "pt",
|
|
).to(model.device)
|
|
prompt_text = input_text
|
|
else:
|
|
# Text-only path for a vision model
|
|
formatted_prompt = self.format_chat_prompt(messages, system_prompt)
|
|
inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to(model.device)
|
|
prompt_text = formatted_prompt
|
|
|
|
# Stream with TextIteratorStreamer + background thread
|
|
try:
|
|
# Re-emit an open <think> prefill swallowed by skip_prompt (see
|
|
# generate_stream).
|
|
think_prefix = detect_think_prefill(
|
|
prompt_text, getattr(raw_tokenizer, "all_special_tokens", None)
|
|
)
|
|
import threading
|
|
|
|
streamer = self._make_text_streamer(
|
|
raw_tokenizer,
|
|
protocol_source = processor,
|
|
# The text-only VLM fallback above did not render with the
|
|
# processor template, so its native markers do not describe
|
|
# this request's response protocol.
|
|
reasoning_channel_markers = detect_reasoning_channel_markers(processor)
|
|
if image
|
|
else None,
|
|
reasoning_channel_markers_resolved = True,
|
|
skip_prompt = True,
|
|
timeout = 0.2,
|
|
cancel_event = cancel_event,
|
|
use_harmony = self._is_gpt_oss_model(),
|
|
)
|
|
|
|
generation_kwargs = dict(
|
|
**inputs,
|
|
streamer = streamer,
|
|
max_new_tokens = max_new_tokens,
|
|
use_cache = True,
|
|
do_sample = temperature > 0,
|
|
temperature = temperature,
|
|
top_p = top_p,
|
|
top_k = top_k,
|
|
min_p = min_p,
|
|
)
|
|
# Presence penalty (GGUF parity) for VLM chat.
|
|
_vision_input_ids = inputs.get("input_ids") if hasattr(inputs, "get") else None
|
|
if _vision_input_ids is not None:
|
|
_pp = _make_presence_penalty_processor(
|
|
presence_penalty, int(_vision_input_ids.shape[1])
|
|
)
|
|
if _pp is not None:
|
|
generation_kwargs["logits_processor"] = _pp
|
|
stopping_criteria = self._cancel_stopping_criteria(cancel_event)
|
|
if stopping_criteria is not None:
|
|
generation_kwargs["stopping_criteria"] = stopping_criteria
|
|
active_stop_token_ids = self._generation_stop_token_ids(model, generation_kwargs)
|
|
|
|
err: dict[str, str] = {}
|
|
|
|
def generate_fn():
|
|
with self._generation_lock:
|
|
try:
|
|
model.generate(**generation_kwargs)
|
|
except Exception as e:
|
|
err["msg"] = str(e)
|
|
if hasattr(streamer, "abort"):
|
|
streamer.abort()
|
|
logger.error(f"Vision generation error in thread: {e}")
|
|
finally:
|
|
try:
|
|
streamer.end()
|
|
except Exception:
|
|
pass
|
|
|
|
thread = threading.Thread(target = generate_fn)
|
|
thread.start()
|
|
|
|
output = think_prefix
|
|
# Emit the prefilled <think> before the first token so the block
|
|
# renders during prompt prefill (which can take seconds).
|
|
if think_prefix:
|
|
yield think_prefix
|
|
from queue import Empty
|
|
import time
|
|
|
|
generation_complete = False
|
|
cancel_deadline = None
|
|
try:
|
|
while True:
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
if cancel_deadline is None:
|
|
cancel_deadline = time.monotonic() + 10
|
|
elif time.monotonic() >= cancel_deadline:
|
|
break
|
|
try:
|
|
new_token = next(streamer)
|
|
except StopIteration:
|
|
generation_complete = True
|
|
break
|
|
except Empty:
|
|
if not thread.is_alive():
|
|
generation_complete = True
|
|
output = yield from self._drain_streamer_tail(
|
|
streamer, output, active_stop_token_ids
|
|
)
|
|
break
|
|
if cancel_deadline is not None:
|
|
remaining = cancel_deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
break
|
|
thread.join(timeout = remaining)
|
|
if thread.is_alive():
|
|
break
|
|
generation_complete = True
|
|
output = yield from self._drain_streamer_tail(
|
|
streamer, output, active_stop_token_ids
|
|
)
|
|
break
|
|
continue
|
|
if new_token:
|
|
output, cleaned = self._append_stream_delta(
|
|
output, new_token, active_stop_token_ids
|
|
)
|
|
yield cleaned
|
|
finally:
|
|
if cancel_event is not None and not generation_complete:
|
|
cancel_event.set()
|
|
join_timeout = 10
|
|
if cancel_deadline is not None:
|
|
join_timeout = max(0, cancel_deadline - time.monotonic())
|
|
thread.join(timeout = join_timeout)
|
|
if thread.is_alive():
|
|
logger.warning(
|
|
"Vision generation thread did not exit after cancel/join timeout"
|
|
)
|
|
|
|
if err.get("msg"):
|
|
raise _GenerationThreadError(err["msg"])
|
|
|
|
except _GenerationThreadError:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Vision generation error: {e}")
|
|
raise
|
|
|
|
def generate_audio_input_response(
|
|
self,
|
|
messages,
|
|
system_prompt,
|
|
audio_array,
|
|
temperature,
|
|
top_p,
|
|
top_k,
|
|
min_p,
|
|
max_new_tokens,
|
|
repetition_penalty,
|
|
cancel_event = None,
|
|
) -> Generator[str, None, None]:
|
|
"""Audio-input (ASR) generation: takes an audio numpy array, streams text.
|
|
|
|
Uses processor.apply_chat_template with audio embedded in messages (Gemma 3n pattern).
|
|
"""
|
|
import threading
|
|
import numpy as np
|
|
|
|
model_info = self.models[self.active_model_name]
|
|
model = model_info["model"]
|
|
processor = model_info.get("processor") or model_info.get("tokenizer")
|
|
raw_tokenizer = getattr(processor, "tokenizer", processor)
|
|
|
|
# Last user text; default matches the notebook prompt
|
|
user_text = "Please transcribe this audio."
|
|
if messages:
|
|
for msg in reversed(messages):
|
|
if msg["role"] == "user" and msg.get("content"):
|
|
user_text = content_to_text(msg["content"])
|
|
break
|
|
|
|
# ASR-specific default system prompt if none set
|
|
if not system_prompt:
|
|
system_prompt = "You are an assistant that transcribes speech accurately."
|
|
|
|
# Gemma 3n format — audio goes INTO apply_chat_template
|
|
audio_messages = [
|
|
{"role": "system", "content": [{"type": "text", "text": system_prompt}]},
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "audio", "audio": audio_array},
|
|
{"type": "text", "text": user_text},
|
|
],
|
|
},
|
|
]
|
|
|
|
# apply_chat_template does audio embedding + tokenization in one step
|
|
inputs = processor.apply_chat_template(
|
|
audio_messages,
|
|
add_generation_prompt = True,
|
|
tokenize = True,
|
|
return_dict = True,
|
|
return_tensors = "pt",
|
|
truncation = False,
|
|
).to(model.device)
|
|
|
|
try:
|
|
from transformers import TextIteratorStreamer
|
|
from queue import Empty
|
|
|
|
streamer = TextIteratorStreamer(
|
|
raw_tokenizer,
|
|
skip_prompt = True,
|
|
skip_special_tokens = True,
|
|
timeout = 0.2,
|
|
)
|
|
|
|
# Notebook uses do_sample=False (greedy) for ASR accuracy
|
|
generation_kwargs = dict(
|
|
**inputs,
|
|
streamer = streamer,
|
|
max_new_tokens = max_new_tokens,
|
|
use_cache = True,
|
|
do_sample = False,
|
|
)
|
|
|
|
err: dict[str, str] = {}
|
|
|
|
def generate_fn():
|
|
with self._generation_lock:
|
|
try:
|
|
model.generate(**generation_kwargs)
|
|
except Exception as e:
|
|
err["msg"] = str(e)
|
|
logger.error(f"Audio input generation error in thread: {e}")
|
|
finally:
|
|
try:
|
|
streamer.end()
|
|
except Exception:
|
|
pass
|
|
|
|
thread = threading.Thread(target = generate_fn)
|
|
thread.start()
|
|
|
|
output = ""
|
|
try:
|
|
while True:
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
break
|
|
try:
|
|
new_token = next(streamer)
|
|
except StopIteration:
|
|
break
|
|
except Empty:
|
|
if not thread.is_alive():
|
|
break
|
|
continue
|
|
if new_token:
|
|
output += new_token
|
|
yield new_token
|
|
finally:
|
|
if cancel_event is not None:
|
|
cancel_event.set()
|
|
thread.join(timeout = 10)
|
|
if thread.is_alive():
|
|
logger.warning(
|
|
"Audio input generation thread did not exit after cancel/join timeout"
|
|
)
|
|
|
|
if err.get("msg"):
|
|
raise _GenerationThreadError(err["msg"])
|
|
|
|
except _GenerationThreadError:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Audio input generation error: {e}")
|
|
raise
|
|
|
|
def generate_whisper_response(
|
|
self,
|
|
audio_array,
|
|
cancel_event = None,
|
|
) -> Generator[str, None, None]:
|
|
"""Whisper ASR: takes an audio numpy array, yields transcribed text.
|
|
|
|
Uses the pre-built transformers pipeline created at model load.
|
|
"""
|
|
model_info = self.models[self.active_model_name]
|
|
whisper_pipe = model_info.get("whisper_pipeline")
|
|
if not whisper_pipe:
|
|
yield "Error: Whisper pipeline not initialized"
|
|
return
|
|
|
|
try:
|
|
with self._generation_lock:
|
|
result = whisper_pipe({"raw": audio_array, "sampling_rate": 16000})
|
|
|
|
text = result.get("text", "") if isinstance(result, dict) else str(result)
|
|
if text:
|
|
yield text
|
|
except Exception as e:
|
|
logger.error(f"Whisper ASR error: {e}")
|
|
yield f"Error: {str(e)}"
|
|
|
|
def _is_gpt_oss_model(self, model_name: str = None) -> bool:
|
|
"""Whether the given (or active) model uses the gpt-oss harmony protocol."""
|
|
from utils.datasets import is_gpt_oss_model_name
|
|
return is_gpt_oss_model_name(model_name or self.active_model_name or "")
|
|
|
|
def _make_text_streamer(
|
|
self,
|
|
tokenizer,
|
|
*,
|
|
protocol_source = None,
|
|
reasoning_channel_markers = None,
|
|
reasoning_channel_markers_resolved: bool = False,
|
|
skip_prompt: bool = True,
|
|
timeout: float = 0.2,
|
|
cancel_event = None,
|
|
use_harmony: bool = False,
|
|
):
|
|
"""Create the streamer matching this model's native response protocol."""
|
|
if use_harmony:
|
|
try:
|
|
return HarmonyTextStreamer(
|
|
tokenizer,
|
|
skip_prompt = skip_prompt,
|
|
timeout = timeout,
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"HarmonyTextStreamer init failed, falling back: {e}")
|
|
return TextIteratorStreamer(
|
|
tokenizer,
|
|
skip_prompt = skip_prompt,
|
|
skip_special_tokens = True,
|
|
timeout = timeout,
|
|
)
|
|
|
|
markers = (
|
|
reasoning_channel_markers
|
|
if reasoning_channel_markers_resolved
|
|
else reasoning_channel_markers
|
|
or detect_reasoning_channel_markers(protocol_source or tokenizer)
|
|
)
|
|
if markers is not None:
|
|
return ReasoningTextIteratorStreamer(
|
|
tokenizer,
|
|
markers = markers,
|
|
skip_prompt = skip_prompt,
|
|
timeout = timeout,
|
|
cancel_event = cancel_event,
|
|
)
|
|
return TextIteratorStreamer(
|
|
tokenizer,
|
|
skip_prompt = skip_prompt,
|
|
skip_special_tokens = True,
|
|
timeout = timeout,
|
|
)
|
|
|
|
def _append_stream_delta(
|
|
self,
|
|
output: str,
|
|
new_token: str,
|
|
stop_token_ids = None,
|
|
):
|
|
"""Append a streamer delta and apply response-boundary cleanup."""
|
|
output += new_token
|
|
return output, self._clean_generated_text(output, stop_token_ids = stop_token_ids)
|
|
|
|
def _drain_streamer_tail(
|
|
self,
|
|
streamer,
|
|
output: str,
|
|
stop_token_ids = None,
|
|
):
|
|
"""Drain queued streamer text after the producer exits."""
|
|
while True:
|
|
try:
|
|
new_token = next(streamer)
|
|
except StopIteration:
|
|
return output
|
|
except Exception:
|
|
return output
|
|
if new_token:
|
|
output, cleaned = self._append_stream_delta(
|
|
output, new_token, stop_token_ids = stop_token_ids
|
|
)
|
|
yield cleaned
|
|
|
|
def generate_stream(
|
|
self,
|
|
prompt: str,
|
|
temperature: float = 0.7,
|
|
top_p: float = 0.9,
|
|
top_k: int = 40,
|
|
min_p: float = 0.0,
|
|
max_new_tokens: int = 256,
|
|
repetition_penalty: float = 1.0,
|
|
cancel_event = None,
|
|
_adapter_state = None,
|
|
presence_penalty: float = 0.0,
|
|
reasoning_channel_markers = None,
|
|
reasoning_channel_markers_resolved: bool = False,
|
|
) -> Generator[str, None, None]:
|
|
"""Generate a streaming text response (text models only).
|
|
|
|
_adapter_state: if not None, the background thread toggles adapters
|
|
before model.generate(), under _generation_lock.
|
|
``presence_penalty`` matches the GGUF sampling path via a logits processor (0 disables it).
|
|
"""
|
|
if not self.active_model_name:
|
|
raise RuntimeError("No active model")
|
|
|
|
model_info = self.models[self.active_model_name]
|
|
model = model_info["model"]
|
|
# For VLMs the stored "tokenizer" is actually the processor. Unwrap to
|
|
# the real tokenizer so TextIteratorStreamer's skip_prompt /
|
|
# skip_special_tokens work correctly.
|
|
tokenizer = model_info["tokenizer"]
|
|
tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
|
|
|
|
try:
|
|
inputs = tokenizer(prompt, return_tensors = "pt").to(model.device)
|
|
|
|
import threading
|
|
|
|
# skip_prompt swallows an open <think> prefilled by the template;
|
|
# re-emit it so the frontend can render the thinking block.
|
|
# gpt-oss emits its own tags via HarmonyTextStreamer.
|
|
think_prefix = (
|
|
""
|
|
if self._is_gpt_oss_model()
|
|
else detect_think_prefill(prompt, getattr(tokenizer, "all_special_tokens", None))
|
|
)
|
|
|
|
streamer = self._make_text_streamer(
|
|
tokenizer,
|
|
protocol_source = model_info.get("tokenizer"),
|
|
reasoning_channel_markers = reasoning_channel_markers,
|
|
reasoning_channel_markers_resolved = reasoning_channel_markers_resolved,
|
|
skip_prompt = True,
|
|
timeout = 0.2,
|
|
cancel_event = cancel_event,
|
|
use_harmony = self._is_gpt_oss_model(),
|
|
)
|
|
|
|
generation_kwargs = dict(
|
|
**inputs,
|
|
streamer = streamer,
|
|
max_new_tokens = max_new_tokens,
|
|
temperature = temperature,
|
|
top_p = top_p,
|
|
top_k = top_k,
|
|
min_p = min_p,
|
|
repetition_penalty = repetition_penalty,
|
|
do_sample = temperature > 0,
|
|
# Resolved once at load (chat_template-derived turn-end tokens).
|
|
eos_token_id = model_info.get("chat_turn_end_eos_ids") or tokenizer.eos_token_id,
|
|
pad_token_id = tokenizer.eos_token_id
|
|
if tokenizer.pad_token_id is None
|
|
else tokenizer.pad_token_id,
|
|
)
|
|
active_stop_token_ids = self._generation_stop_token_ids(model, generation_kwargs)
|
|
# Presence penalty (GGUF parity); prompt_len excludes prompt tokens.
|
|
_pp = _make_presence_penalty_processor(
|
|
presence_penalty, int(inputs["input_ids"].shape[1])
|
|
)
|
|
if _pp is not None:
|
|
generation_kwargs["logits_processor"] = _pp
|
|
stopping_criteria = self._cancel_stopping_criteria(cancel_event)
|
|
if stopping_criteria is not None:
|
|
generation_kwargs["stopping_criteria"] = stopping_criteria
|
|
|
|
def generate_fn():
|
|
with self._generation_lock:
|
|
try:
|
|
if _adapter_state is not None:
|
|
self._apply_adapter_state(_adapter_state)
|
|
model.generate(**generation_kwargs)
|
|
except Exception as e:
|
|
err["msg"] = str(e)
|
|
if hasattr(streamer, "abort"):
|
|
streamer.abort()
|
|
logger.error(f"Generation error: {e}")
|
|
finally:
|
|
try:
|
|
streamer.end()
|
|
except Exception:
|
|
pass
|
|
|
|
err: dict[str, str] = {}
|
|
thread = threading.Thread(target = generate_fn)
|
|
thread.start()
|
|
|
|
output = think_prefix
|
|
# Emit the prefilled <think> before the first token so the block
|
|
# renders during prompt prefill (which can take seconds).
|
|
if think_prefix:
|
|
yield think_prefix
|
|
from queue import Empty
|
|
import time
|
|
|
|
generation_complete = False
|
|
cancel_deadline = None
|
|
try:
|
|
while True:
|
|
if cancel_event is not None and cancel_event.is_set():
|
|
if cancel_deadline is None:
|
|
cancel_deadline = time.monotonic() + 10
|
|
elif time.monotonic() >= cancel_deadline:
|
|
break
|
|
try:
|
|
new_token = next(streamer)
|
|
except StopIteration:
|
|
generation_complete = True
|
|
break
|
|
except Empty:
|
|
if not thread.is_alive():
|
|
generation_complete = True
|
|
output = yield from self._drain_streamer_tail(
|
|
streamer, output, active_stop_token_ids
|
|
)
|
|
break
|
|
if cancel_deadline is not None:
|
|
remaining = cancel_deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
break
|
|
thread.join(timeout = remaining)
|
|
if thread.is_alive():
|
|
break
|
|
generation_complete = True
|
|
output = yield from self._drain_streamer_tail(
|
|
streamer, output, active_stop_token_ids
|
|
)
|
|
break
|
|
continue
|
|
if new_token:
|
|
output, cleaned = self._append_stream_delta(
|
|
output, new_token, active_stop_token_ids
|
|
)
|
|
yield cleaned
|
|
finally:
|
|
# Set cancel_event only on early exit (user cancel), NOT on
|
|
# normal completion. It's a shared mp.Event; setting it
|
|
# unconditionally would leave a stale cancel signal that could
|
|
# disrupt the next serialized request (e.g. compare mode).
|
|
if cancel_event is not None and not generation_complete:
|
|
cancel_event.set()
|
|
join_timeout = 10
|
|
if cancel_deadline is not None:
|
|
join_timeout = max(0, cancel_deadline - time.monotonic())
|
|
thread.join(timeout = join_timeout)
|
|
if thread.is_alive():
|
|
logger.warning("Generation thread did not exit after cancel/join timeout")
|
|
|
|
if err.get("msg"):
|
|
raise _GenerationThreadError(err["msg"])
|
|
|
|
except _GenerationThreadError:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Error during generation: {e}")
|
|
raise
|
|
|
|
# ── Audio (TTS) Generation ────────────────────────────────────
|
|
|
|
def generate_audio_response(
|
|
self,
|
|
text: str,
|
|
temperature: float = 0.6,
|
|
top_p: float = 0.95,
|
|
top_k: int = 50,
|
|
min_p: float = 0.0,
|
|
max_new_tokens: int = 2048,
|
|
repetition_penalty: float = 1.0,
|
|
use_adapter: Optional[Union[bool, str]] = None,
|
|
) -> Tuple[bytes, int]:
|
|
"""Generate audio from text for TTS models.
|
|
Returns (wav_bytes, sample_rate). Blocking — full audio before return.
|
|
"""
|
|
if not self.active_model_name:
|
|
raise RuntimeError("No active model")
|
|
|
|
model_info = self.models[self.active_model_name]
|
|
audio_type = model_info.get("audio_type")
|
|
model = model_info["model"]
|
|
tokenizer = model_info.get("tokenizer")
|
|
|
|
if not audio_type:
|
|
raise RuntimeError(f"Model {self.active_model_name} is not an audio model")
|
|
|
|
top_k = self._normalize_top_k(top_k)
|
|
|
|
with self._generation_lock:
|
|
if use_adapter is not None:
|
|
self._apply_adapter_state(use_adapter)
|
|
|
|
if audio_type == "snac":
|
|
return self._generate_snac(
|
|
model,
|
|
tokenizer,
|
|
text,
|
|
temperature,
|
|
top_p,
|
|
max_new_tokens,
|
|
repetition_penalty,
|
|
)
|
|
elif audio_type == "csm":
|
|
processor = model_info.get("processor", tokenizer)
|
|
return self._generate_csm(model, processor, text, max_new_tokens)
|
|
elif audio_type == "bicodec":
|
|
return self._generate_bicodec(
|
|
model, tokenizer, text, temperature, top_k, max_new_tokens
|
|
)
|
|
elif audio_type == "dac":
|
|
return self._generate_dac(
|
|
model,
|
|
tokenizer,
|
|
text,
|
|
temperature,
|
|
top_k,
|
|
top_p,
|
|
min_p,
|
|
max_new_tokens,
|
|
repetition_penalty,
|
|
)
|
|
else:
|
|
raise RuntimeError(f"Unknown audio_type: {audio_type}")
|
|
|
|
def _generate_snac(
|
|
self, model, tokenizer, text, temperature, top_p, max_new_tokens, repetition_penalty
|
|
):
|
|
"""Generate audio using SNAC codec (Orpheus)."""
|
|
device = model.device
|
|
start_token = torch.tensor([[128259]], device = device) # START_OF_HUMAN
|
|
end_tokens = torch.tensor([[128009, 128260]], device = device) # EOT, END_OF_HUMAN
|
|
text_ids = tokenizer(text, return_tensors = "pt").input_ids.to(device)
|
|
input_ids = torch.cat([start_token, text_ids, end_tokens], dim = 1)
|
|
attention_mask = torch.ones_like(input_ids)
|
|
|
|
generated = model.generate(
|
|
input_ids = input_ids,
|
|
attention_mask = attention_mask,
|
|
max_new_tokens = max_new_tokens,
|
|
do_sample = True,
|
|
temperature = temperature,
|
|
top_p = top_p,
|
|
repetition_penalty = repetition_penalty,
|
|
eos_token_id = 128258, # END_OF_SPEECH
|
|
use_cache = True,
|
|
)
|
|
return self._audio_codec_manager.decode_snac(generated, str(device))
|
|
|
|
def _generate_csm(self, model, processor, text, max_new_tokens):
|
|
"""Generate audio using CSM (Sesame)."""
|
|
speaker_id = 0
|
|
inputs = processor(
|
|
f"[{speaker_id}]{text}", add_special_tokens = True, return_tensors = "pt"
|
|
).to(model.device)
|
|
audio_values = model.generate(**inputs, max_new_tokens = max_new_tokens, output_audio = True)
|
|
return self._audio_codec_manager.decode_csm(audio_values)
|
|
|
|
def _generate_bicodec(self, model, tokenizer, text, temperature, top_k, max_new_tokens):
|
|
"""Generate audio using BiCodec (Spark-TTS)."""
|
|
prompt = "<|task_tts|><|start_content|>" + text + "<|end_content|><|start_global_token|>"
|
|
inputs = tokenizer([prompt], return_tensors = "pt").to(model.device)
|
|
generated = model.generate(
|
|
**inputs,
|
|
max_new_tokens = max_new_tokens,
|
|
do_sample = True,
|
|
temperature = temperature,
|
|
top_k = top_k,
|
|
eos_token_id = tokenizer.eos_token_id,
|
|
pad_token_id = tokenizer.pad_token_id,
|
|
)
|
|
new_tokens = generated[:, inputs.input_ids.shape[1] :]
|
|
decoded_text = tokenizer.batch_decode(new_tokens, skip_special_tokens = False)[0]
|
|
return self._audio_codec_manager.decode_bicodec(decoded_text, str(model.device))
|
|
|
|
def _generate_dac(
|
|
self,
|
|
model,
|
|
tokenizer,
|
|
text,
|
|
temperature,
|
|
top_k,
|
|
top_p,
|
|
min_p,
|
|
max_new_tokens,
|
|
repetition_penalty,
|
|
):
|
|
"""Generate audio using DAC (OuteTTS). Follows Oute_TTS_(1B).ipynb exactly."""
|
|
# Monkey-patch RepetitionPenaltyLogitsProcessor with a 64-token window
|
|
# (same as the OuteTTS notebook) to avoid degenerate repetition.
|
|
self._patch_repetition_penalty_processor()
|
|
|
|
prompt = (
|
|
"<|im_start|>\n<|text_start|>"
|
|
+ text
|
|
+ "<|text_end|>\n<|audio_start|><|global_features_start|>\n"
|
|
)
|
|
|
|
with torch.inference_mode():
|
|
# Derive the autocast device from the loaded model, not from the
|
|
# global backend: a CPU-fallback DAC on an XPU/CUDA host must not
|
|
# open a GPU autocast context around CPU tensors.
|
|
device_type = (
|
|
model.device.type
|
|
if hasattr(model.device, "type")
|
|
else str(model.device).split(":", 1)[0]
|
|
)
|
|
# Clamp to autocast-supported backends so exotic devices
|
|
# (e.g. "meta" during accelerate offloaded loading) do not raise.
|
|
# MPS is autocast-supported since torch 2.3, keep it in the set.
|
|
if device_type not in ("cuda", "xpu", "mps", "cpu"):
|
|
device_type = "cpu"
|
|
# CPU and XPU autocast only accept bfloat16/float16. For a
|
|
# float32 model, skip autocast entirely to avoid raising or
|
|
# producing a warning on every generate call.
|
|
autocast_dtype_supported = model.dtype in (torch.bfloat16, torch.float16)
|
|
if device_type in ("cpu", "xpu") and not autocast_dtype_supported:
|
|
autocast_ctx = contextlib.nullcontext()
|
|
else:
|
|
autocast_ctx = torch.amp.autocast(device_type, dtype = model.dtype)
|
|
with autocast_ctx:
|
|
inputs = tokenizer([prompt], return_tensors = "pt").to(model.device)
|
|
generated = model.generate(
|
|
**inputs,
|
|
temperature = temperature,
|
|
top_k = top_k,
|
|
top_p = top_p,
|
|
min_p = min_p,
|
|
repetition_penalty = repetition_penalty,
|
|
max_new_tokens = max_new_tokens,
|
|
)
|
|
decoded_text = tokenizer.batch_decode(generated, skip_special_tokens = False)[0]
|
|
return self._audio_codec_manager.decode_dac(decoded_text, str(model.device))
|
|
|
|
_repetition_penalty_patched = False
|
|
|
|
@classmethod
|
|
def _patch_repetition_penalty_processor(cls):
|
|
"""Monkey-patch transformers' RepetitionPenaltyLogitsProcessor with a
|
|
64-token sliding-window variant (from the OuteTTS notebook).
|
|
Applied once per process.
|
|
"""
|
|
if cls._repetition_penalty_patched:
|
|
return
|
|
cls._repetition_penalty_patched = True
|
|
|
|
from transformers import LogitsProcessor
|
|
import transformers.generation.utils as generation_utils
|
|
|
|
class RepetitionPenaltyLogitsProcessorPatch(LogitsProcessor):
|
|
def __init__(self, penalty: float):
|
|
self.penalty_last_n = 64
|
|
if not isinstance(penalty, float) or penalty <= 0:
|
|
raise ValueError(f"`penalty` has to be a positive float, but is {penalty}")
|
|
self.penalty = penalty
|
|
|
|
@torch.no_grad()
|
|
def __call__(
|
|
self, input_ids: torch.LongTensor, scores: torch.FloatTensor
|
|
) -> torch.FloatTensor:
|
|
if self.penalty_last_n == 0 or self.penalty == 1.0:
|
|
return scores
|
|
batch_size, seq_len = input_ids.shape
|
|
vocab_size = scores.shape[-1]
|
|
for b in range(batch_size):
|
|
start_index = max(0, seq_len - self.penalty_last_n)
|
|
window_indices = input_ids[b, start_index:]
|
|
if window_indices.numel() == 0:
|
|
continue
|
|
for token_id in set(window_indices.tolist()):
|
|
if token_id >= vocab_size:
|
|
continue
|
|
logit = scores[b, token_id]
|
|
scores[b, token_id] = (
|
|
logit * self.penalty if logit <= 0 else logit / self.penalty
|
|
)
|
|
return scores
|
|
|
|
generation_utils.RepetitionPenaltyLogitsProcessor = RepetitionPenaltyLogitsProcessorPatch
|
|
logger.info("Patched RepetitionPenaltyLogitsProcessor with 64-token window for OuteTTS")
|
|
|
|
def _apply_chat_template_for_generation(
|
|
self,
|
|
tokenizer,
|
|
messages: list,
|
|
*,
|
|
tools: Optional[list] = None,
|
|
enable_thinking: Optional[bool] = None,
|
|
reasoning_effort: Optional[str] = None,
|
|
preserve_thinking: Optional[bool] = None,
|
|
) -> str:
|
|
"""Render the chat prompt, peeling kwargs the template doesn't
|
|
understand. Delegates to the dependency-light helper module so the
|
|
fallback chain is unit-testable without pulling unsloth / torch into
|
|
the test sandbox.
|
|
"""
|
|
from core.inference.chat_template_helpers import (
|
|
apply_chat_template_for_generation,
|
|
)
|
|
return apply_chat_template_for_generation(
|
|
tokenizer,
|
|
messages,
|
|
tools = tools,
|
|
enable_thinking = enable_thinking,
|
|
reasoning_effort = reasoning_effort,
|
|
preserve_thinking = preserve_thinking,
|
|
)
|
|
|
|
def format_chat_prompt(
|
|
self,
|
|
messages: list,
|
|
system_prompt: str = None,
|
|
) -> str:
|
|
if not self.active_model_name or self.active_model_name not in self.models:
|
|
logger.error("No active model available")
|
|
return ""
|
|
|
|
if self.models[self.active_model_name].get("tokenizer") is None:
|
|
logger.error("Tokenizer not loaded for active model")
|
|
return ""
|
|
|
|
chat_template_info = self.models[self.active_model_name].get("chat_template_info", {})
|
|
tokenizer = self.models[self.active_model_name]["tokenizer"]
|
|
tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
|
|
|
|
chat_messages = []
|
|
|
|
if system_prompt:
|
|
chat_messages.append({"role": "system", "content": system_prompt})
|
|
|
|
last_role = "system" if system_prompt else None
|
|
|
|
for msg in messages:
|
|
role = msg.get("role", "")
|
|
content = content_to_text(msg.get("content", ""))
|
|
|
|
if role in ["system", "user", "assistant"] and content.strip():
|
|
if role == last_role:
|
|
logger.debug(f"Skipping consecutive {role} message to maintain alternation")
|
|
continue
|
|
|
|
if role == "user":
|
|
import re
|
|
clean_content = re.sub(r"<[^>]+>", "", content).strip()
|
|
if clean_content:
|
|
chat_messages.append({"role": role, "content": clean_content})
|
|
last_role = role
|
|
elif role == "assistant" and content.strip():
|
|
chat_messages.append({"role": role, "content": content})
|
|
last_role = role
|
|
elif role == "system":
|
|
continue
|
|
|
|
if chat_messages and chat_messages[-1]["role"] == "assistant":
|
|
logger.debug("Removing final assistant message to ensure proper alternation")
|
|
chat_messages.pop()
|
|
|
|
logger.info(f"Sending {len(chat_messages)} messages to tokenizer:")
|
|
for i, msg in enumerate(chat_messages):
|
|
logger.info(f" {i}: {msg['role']} - {msg['content'][:50]}...")
|
|
|
|
try:
|
|
formatted_prompt = tokenizer.apply_chat_template(
|
|
chat_messages, tokenize = False, add_generation_prompt = True
|
|
)
|
|
logger.info(f"Successfully applied tokenizer's native chat template")
|
|
return formatted_prompt
|
|
except Exception as e:
|
|
error_msg = str(e).lower()
|
|
if "chat_template is not set" in error_msg or "no template argument" in error_msg:
|
|
logger.info(
|
|
f"Base model detected - no built-in chat template available, using fallback formatting"
|
|
)
|
|
else:
|
|
logger.warning(f"Failed to apply tokenizer chat template: {e}")
|
|
logger.debug(
|
|
f"""Failed with messages: {[f"{m['role']}: {m['content'][:30]}..." for m in chat_messages]}"""
|
|
)
|
|
|
|
if chat_template_info.get("has_template", False):
|
|
logger.info("Falling back to manual template formatting based on detected patterns")
|
|
template_type = chat_template_info.get("format_type", "generic")
|
|
manual_prompt = self._format_chat_manual(
|
|
chat_messages,
|
|
template_type,
|
|
chat_template_info.get("special_tokens", {}),
|
|
)
|
|
logger.info(f"Manual template result: {manual_prompt[:200]}...")
|
|
return manual_prompt
|
|
else:
|
|
logger.info("Using generic chat formatting for base model")
|
|
return self._format_generic_template(chat_messages, {})
|
|
|
|
def _format_chat_manual(self, messages: list, template_type: str, special_tokens: dict) -> str:
|
|
"""Manual chat-formatting fallback when the tokenizer template fails.
|
|
|
|
Args:
|
|
messages: List of message dictionaries
|
|
template_type: Detected template type
|
|
special_tokens: Dictionary of special tokens
|
|
|
|
Returns:
|
|
str: Manually formatted prompt
|
|
"""
|
|
if template_type == "llama3":
|
|
return self._format_llama3_template(messages, special_tokens)
|
|
elif template_type == "mistral":
|
|
return self._format_mistral_template(messages, special_tokens)
|
|
elif template_type == "chatml":
|
|
return self._format_chatml_template(messages, special_tokens)
|
|
elif template_type == "alpaca":
|
|
return self._format_alpaca_template(messages, special_tokens)
|
|
else:
|
|
return self._format_generic_template(messages, special_tokens)
|
|
|
|
def _format_llama3_template(self, messages: list, special_tokens: dict) -> str:
|
|
"""Format messages using Llama 3 template"""
|
|
bos_token = special_tokens.get("bos_token", "<|begin_of_text|>")
|
|
formatted = bos_token
|
|
|
|
for msg in messages:
|
|
role = msg["role"]
|
|
content = content_to_text(msg["content"])
|
|
formatted += f"<|start_header_id|>{role}<|end_header_id|>\n\n{content}<|eot_id|>"
|
|
|
|
formatted += "<|start_header_id|>assistant<|end_header_id|>\n\n"
|
|
return formatted
|
|
|
|
def _format_mistral_template(self, messages: list, special_tokens: dict) -> str:
|
|
"""Format messages using Mistral template"""
|
|
bos_token = special_tokens.get("bos_token", "<s>")
|
|
formatted = bos_token
|
|
|
|
system_msg = None
|
|
conversation = []
|
|
|
|
for msg in messages:
|
|
if msg["role"] == "system":
|
|
system_msg = content_to_text(msg["content"])
|
|
else:
|
|
conversation.append(msg)
|
|
|
|
i = 0
|
|
while i < len(conversation):
|
|
if conversation[i]["role"] == "user":
|
|
user_content = content_to_text(conversation[i]["content"])
|
|
|
|
if system_msg and i == 0:
|
|
user_content = f"{system_msg}\n\n{user_content}"
|
|
|
|
formatted += f"[INST] {user_content} [/INST]"
|
|
|
|
if i + 1 < len(conversation) and conversation[i + 1]["role"] == "assistant":
|
|
formatted += f" {content_to_text(conversation[i + 1]['content'])}</s>"
|
|
i += 2
|
|
else:
|
|
formatted += " "
|
|
break
|
|
else:
|
|
i += 1
|
|
|
|
return formatted
|
|
|
|
def _format_chatml_template(self, messages: list, special_tokens: dict) -> str:
|
|
"""Format messages using ChatML template"""
|
|
formatted = ""
|
|
|
|
for msg in messages:
|
|
role = msg["role"]
|
|
content = content_to_text(msg["content"])
|
|
formatted += f"<|im_start|>{role}\n{content}<|im_end|>\n"
|
|
|
|
formatted += "<|im_start|>assistant\n"
|
|
return formatted
|
|
|
|
def _format_alpaca_template(self, messages: list, special_tokens: dict) -> str:
|
|
"""Format messages using Alpaca template"""
|
|
formatted = ""
|
|
system_msg = None
|
|
|
|
for msg in messages:
|
|
content = content_to_text(msg["content"])
|
|
if msg["role"] == "system":
|
|
system_msg = content
|
|
elif msg["role"] == "user":
|
|
if system_msg:
|
|
formatted += f"### Instruction:\n{system_msg}\n\n### Input:\n{content}\n\n### Response:\n"
|
|
system_msg = None
|
|
else:
|
|
formatted += f"### Human:\n{content}\n\n### Assistant:\n"
|
|
elif msg["role"] == "assistant":
|
|
formatted += f"{content}\n\n"
|
|
|
|
return formatted
|
|
|
|
def _format_generic_template(self, messages: list, special_tokens: dict) -> str:
|
|
"""Generic fallback formatting"""
|
|
formatted = ""
|
|
|
|
for msg in messages:
|
|
role = msg["role"].title()
|
|
content = content_to_text(msg["content"])
|
|
formatted += f"{role}: {content}\n"
|
|
|
|
formatted += "Assistant: "
|
|
return formatted
|
|
|
|
def check_vision_model_compatibility(self) -> bool:
|
|
"""Whether the current model supports vision."""
|
|
current_model = self.get_current_model()
|
|
if current_model and current_model in self.models:
|
|
return self.models[current_model].get("is_vision", False)
|
|
return False
|
|
|
|
def _reset_model_generation_state(self, model_name: str):
|
|
"""Reset generation state for a specific model to prevent contamination."""
|
|
if model_name not in self.models:
|
|
return
|
|
|
|
model = self.models[model_name].get("model")
|
|
if not model:
|
|
return
|
|
|
|
try:
|
|
# Common pattern for Unsloth/Hugging Face models
|
|
if hasattr(model, "past_key_values"):
|
|
model.past_key_values = None
|
|
if hasattr(model, "generation_config"):
|
|
if hasattr(model.generation_config, "past_key_values"):
|
|
model.generation_config.past_key_values = None
|
|
|
|
logger.debug(f"Reset generation state for model: {model_name}")
|
|
except Exception as e:
|
|
logger.warning(f"Could not fully reset model state for {model_name}: {e}")
|
|
|
|
def reset_generation_state(self, caller_cancel_event = None):
|
|
"""Reset any cached generation state to prevent hanging after errors
|
|
|
|
``caller_cancel_event`` is accepted for signature parity with the
|
|
orchestrator, which uses it to drop a reset from a request that never
|
|
started. Nothing here cancels a live generation, so it is unused.
|
|
"""
|
|
try:
|
|
# Clear cached state for ALL loaded models
|
|
for model_name in self.models.keys():
|
|
self._reset_model_generation_state(model_name)
|
|
|
|
clear_gpu_cache()
|
|
logger.debug("Cleared GPU cache")
|
|
|
|
import gc
|
|
|
|
gc.collect()
|
|
logger.info("Performed comprehensive generation state reset")
|
|
|
|
except Exception as e:
|
|
logger.warning(f"Could not fully reset generation state: {e}")
|
|
|
|
def resize_image(
|
|
self,
|
|
img,
|
|
max_size: int = 800,
|
|
):
|
|
"""Resize image while maintaining aspect ratio if either dimension exceeds max_size"""
|
|
if img is None:
|
|
return None
|
|
if img.size[0] > max_size or img.size[1] > max_size:
|
|
from PIL import Image
|
|
|
|
ratio = min(max_size / img.size[0], max_size / img.size[1])
|
|
new_size = (int(img.size[0] * ratio), int(img.size[1] * ratio))
|
|
return img.resize(new_size, Image.Resampling.LANCZOS)
|
|
return img
|
|
|
|
def _generation_stop_token_ids(self, model, generation_kwargs: dict):
|
|
"""Return the stop-token ids active for a ``generate`` call."""
|
|
if "eos_token_id" in generation_kwargs:
|
|
return generation_kwargs.get("eos_token_id")
|
|
generation_config = getattr(model, "generation_config", None)
|
|
eos_token_id = getattr(generation_config, "eos_token_id", None)
|
|
if eos_token_id is not None:
|
|
return eos_token_id
|
|
config = getattr(model, "config", None)
|
|
return getattr(config, "eos_token_id", None)
|
|
|
|
def _cancel_stopping_criteria(self, cancel_event):
|
|
"""Build a Transformers stopping criteria list for user cancellation."""
|
|
if cancel_event is None:
|
|
return None
|
|
from transformers.generation.stopping_criteria import (
|
|
StoppingCriteria,
|
|
StoppingCriteriaList,
|
|
)
|
|
|
|
class _CancelCriteria(StoppingCriteria):
|
|
def __init__(self, ev):
|
|
self.ev = ev
|
|
|
|
def __call__(self, input_ids, scores, **kwargs):
|
|
return self.ev.is_set()
|
|
|
|
return StoppingCriteriaList([_CancelCriteria(cancel_event)])
|
|
|
|
def _clean_generated_text(
|
|
self,
|
|
text: str,
|
|
*,
|
|
stop_token_ids = None,
|
|
) -> str:
|
|
"""Strip leaked response-boundary tokens after streaming."""
|
|
if self._is_gpt_oss_model():
|
|
# HarmonyTextStreamer emits clean <think>...</think>. Strip any
|
|
# harmony protocol tokens and other gpt-oss tokens (e.g.
|
|
# <|return|>) that leak past the streamer.
|
|
import re
|
|
text = re.sub(r"<\|[a-z_]+\|>", "", text)
|
|
return text.strip()
|
|
|
|
tokenizer = self.models.get(self.active_model_name, {}).get("tokenizer")
|
|
tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
|
|
if tokenizer:
|
|
if stop_token_ids is None:
|
|
stop_token_ids = self.models.get(self.active_model_name, {}).get(
|
|
"chat_turn_end_eos_ids"
|
|
)
|
|
if isinstance(stop_token_ids, int):
|
|
stop_token_ids = (stop_token_ids,)
|
|
for token_id in stop_token_ids or ():
|
|
try:
|
|
token = tokenizer.convert_ids_to_tokens(int(token_id))
|
|
except Exception:
|
|
token = None
|
|
if isinstance(token, str) and token and text.endswith(token):
|
|
text = text[: -len(token)]
|
|
elif (
|
|
isinstance(token, str)
|
|
and token
|
|
and text.endswith("</think>")
|
|
and text[: -len("</think>")].endswith(token)
|
|
):
|
|
text = text[: -len("</think>") - len(token)] + "</think>"
|
|
return text.strip()
|
|
|
|
def _load_chat_template_info(self, model_name: str):
|
|
if model_name not in self.models or not self.models[model_name].get("tokenizer"):
|
|
return
|
|
|
|
tokenizer = self.models[model_name]["tokenizer"]
|
|
chat_template_info = {
|
|
"has_template": False,
|
|
"template": None,
|
|
"format_type": "generic",
|
|
"special_tokens": {},
|
|
"template_name": None,
|
|
}
|
|
|
|
try:
|
|
from utils.datasets import MODEL_TO_TEMPLATE_MAPPER
|
|
|
|
# Exact match first
|
|
model_name_lower = model_name.lower()
|
|
if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
|
|
chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
|
|
logger.info(
|
|
f"Detected template '{chat_template_info['template_name']}' for {model_name} from mapper"
|
|
)
|
|
else:
|
|
# Partial match (for variants like model_name-bnb-4bit)
|
|
for key in MODEL_TO_TEMPLATE_MAPPER:
|
|
if key in model_name_lower or model_name_lower in key:
|
|
chat_template_info["template_name"] = MODEL_TO_TEMPLATE_MAPPER[key]
|
|
logger.info(
|
|
f"Detected template '{chat_template_info['template_name']}' for {model_name} (partial match)"
|
|
)
|
|
break
|
|
except Exception as e:
|
|
logger.warning(f"Could not detect template from mapper for {model_name}: {e}")
|
|
|
|
try:
|
|
if hasattr(tokenizer, "chat_template") and tokenizer.chat_template:
|
|
chat_template_info["has_template"] = True
|
|
chat_template_info["template"] = tokenizer.chat_template
|
|
|
|
template_str = tokenizer.chat_template.lower()
|
|
|
|
if "start_header_id" in template_str and "end_header_id" in template_str:
|
|
chat_template_info["format_type"] = "llama3"
|
|
elif "[inst]" in template_str and "[/inst]" in template_str:
|
|
chat_template_info["format_type"] = "mistral"
|
|
elif "<|im_start|>" in template_str and "<|im_end|>" in template_str:
|
|
chat_template_info["format_type"] = "chatml"
|
|
elif "### instruction:" in template_str or "### human:" in template_str:
|
|
chat_template_info["format_type"] = "alpaca"
|
|
else:
|
|
chat_template_info["format_type"] = "custom"
|
|
|
|
logger.info(
|
|
f"Loaded chat template for {model_name} (detected as {chat_template_info['format_type']} format)"
|
|
)
|
|
logger.debug(f"Template preview: {tokenizer.chat_template[:200]}...")
|
|
|
|
special_tokens = {}
|
|
if hasattr(tokenizer, "bos_token") and tokenizer.bos_token:
|
|
special_tokens["bos_token"] = tokenizer.bos_token
|
|
if hasattr(tokenizer, "eos_token") and tokenizer.eos_token:
|
|
special_tokens["eos_token"] = tokenizer.eos_token
|
|
if hasattr(tokenizer, "pad_token") and tokenizer.pad_token:
|
|
special_tokens["pad_token"] = tokenizer.pad_token
|
|
|
|
chat_template_info["special_tokens"] = special_tokens
|
|
|
|
else:
|
|
logger.info(f"No chat template found for {model_name}, will use generic formatting")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error loading chat template info for {model_name}: {e}")
|
|
|
|
self.models[model_name]["chat_template_info"] = chat_template_info
|
|
|
|
if chat_template_info["has_template"]:
|
|
logger.info(
|
|
f"Chat template loaded for {model_name}: {chat_template_info['format_type']} format"
|
|
)
|
|
else:
|
|
logger.info(f"No built-in chat template for {model_name}, will use generic formatting")
|
|
|
|
def get_current_model(self) -> Optional[str]:
|
|
"""Currently active model name."""
|
|
return self.active_model_name
|
|
|
|
def is_model_loading(self) -> bool:
|
|
"""Whether any model is currently loading."""
|
|
return len(self.loading_models) > 0
|
|
|
|
def get_loading_model(self) -> Optional[str]:
|
|
"""Name of the currently loading model."""
|
|
return next(iter(self.loading_models)) if self.loading_models else None
|
|
|
|
def load_model_simple(
|
|
self,
|
|
model_path: str,
|
|
hf_token: Optional[str] = None,
|
|
max_seq_length: int = 2048,
|
|
load_in_4bit: bool = True,
|
|
) -> bool:
|
|
"""Simple model-loading wrapper for the chat interface. Takes a string
|
|
path and builds the ModelConfig internally.
|
|
|
|
Args:
|
|
model_path: Model name or path (e.g., "unsloth/llama-3-8b")
|
|
hf_token: HuggingFace token for gated models
|
|
max_seq_length: Maximum sequence length
|
|
load_in_4bit: Whether to use 4-bit quantization
|
|
|
|
Returns:
|
|
bool: True if successful, False otherwise
|
|
"""
|
|
try:
|
|
config = ModelConfig.from_ui_selection(
|
|
model_path,
|
|
lora_path = None, # No LoRA for chat
|
|
is_lora = False,
|
|
)
|
|
|
|
return self.load_model(
|
|
config = config,
|
|
max_seq_length = max_seq_length,
|
|
dtype = None, # Auto-detect
|
|
load_in_4bit = load_in_4bit,
|
|
hf_token = hf_token,
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in load_model_simple: {e}")
|
|
return False
|
|
|
|
|
|
# Global inference backend instance
|
|
inference_backend = InferenceBackend()
|
|
|
|
|
|
def get_inference_backend() -> InferenceBackend:
|
|
return inference_backend
|