Bring Your Own Model¶
You trained a model outside mixle entirely – a HuggingFace checkpoint, a
TRL-fine-tuned chat model, an arbitrary predict(x) function from another
framework – and you do not want to rewrite it in mixle’s terms. This page is
the on-ramp: wrap it as a callable teacher, then compose it with mixle’s
distillation, calibration, and routing machinery unchanged.
The shape of the whole page is one pipeline:
externally-trained model -> callable teacher -> distill -> calibrate -> route
Each arrow is one mixle call. Nothing here requires the external model’s weights, framework, or training loop to know anything about mixle.
This is a different interop pattern than Neural and LLM Models’s “drop a real
torch module into GradLeaf” story – see This vs. The GradLeaf Pattern below for
when to reach for which.
1. Wrap It As A Callable Teacher¶
mixle.task.llm.CallableLLM wraps any fn(prompt) -> str (or
fn(prompt, system) -> str) as an LLM. The function can be a thin
shim around model.generate() for an HF checkpoint, a TRL-trained policy,
a llama.cpp binding, or a test double – mixle only ever calls
complete(prompt).
from mixle.task import CallableLLM, llm_labeler
# Stand-in for `tokenizer(prompt) -> model.generate(...) -> tokenizer.decode(...)`
# against a checkpoint trained outside mixle (HF Trainer, TRL SFT/DPO, ...).
def generate(prompt: str, system: str | None = None) -> str:
# e.g. hf_pipeline(prompt)[0]["generated_text"], or a TRL-trained model's .generate()
return "spam" if "free prize" in prompt.lower() else "ham"
teacher_llm = CallableLLM(generate)
# Constrain free-text replies to a fixed label set -- the batched-teacher
# shape (`texts -> [label]`) every mixle.task distillation entry point expects.
teacher = llm_labeler(
teacher_llm,
["spam", "ham"],
instruction="Classify the email as spam or ham.",
)
If the checkpoint is served behind an OpenAI-compatible endpoint (vLLM, TGI,
Ollama, a hosted API) instead of loaded in-process, use OpenAICompatLLM
in place of CallableLLM – same downstream contract, no local inference
code at all.
For a vision-language checkpoint, the sibling is mixle.task.vlm:
CallableVLM wraps fn(image, prefix) -> [(token, log_prob), ...], and
OpenAICompatVLM talks to a logprob-serving OpenAI-compatible endpoint.
Both plug into mixle.enumeration.best_first_decode the same way an LLM
teacher plugs into distillation.
At this point, teacher is indistinguishable to the rest of mixle.task
from a frontier API model – everything below is unaware the labels came
from an HF checkpoint.
2. Distill: Train A Small mixle-Native Student¶
distill exercises the teacher once over an unlabeled corpus and fits a
small local classifier that reproduces its labels:
from mixle.task import distill
student = distill(
teacher,
unlabeled_emails,
labels=["spam", "ham"],
task="spam classifier distilled from an HF/TRL teacher",
)
student.meta["train_agreement"] # how faithfully the student mimics the teacher
student.save("spam_student.mixle")
If labeling budget is the constraint rather than compute, active_distill
spends a fixed label budget on the most informative pool examples
(by margin/entropy acquisition) instead of labeling everything:
from mixle.task import active_distill
result = active_distill(teacher, unlabeled_pool, budget=200, rounds=5)
student = result.model
Either way the teacher is opaque: distill/active_distill only ever
call it as teacher(list_of_texts) -> list_of_labels. Nothing about the
student’s training loop touches the checkpoint, its framework, or its
weights.
3. Calibrate: An Honest Answer-Or-Escalate Guarantee¶
A distilled student’s softmax is not a probability guarantee on its own.
CalibratedTaskModel wraps the student in a conformal threshold so
“answer locally” vs. “escalate” is backed by coverage, not vibes:
from mixle.task import CalibratedTaskModel
calibrated = CalibratedTaskModel(student, alpha=0.1)
calibrated.calibrate(held_out_texts, held_out_labels)
calibrated.decide(new_email) # a label, or ESCALATE (None) if unsure
calibrated.escalation_rate(val_texts) # the empirical p(escalate)
What that guarantee is, exactly – and what it is not. calibrate() gives the
conformal statement: if calibration rows and future queries are exchangeable draws from the
same distribution, the prediction set contains the true label with probability at least
1 - alpha – marginally, over the whole query population. It is NOT a conditional
guarantee on the answered slice: the error rate among the queries the model chooses to
answer can be far higher (an adversarial review measured 9% marginal miscoverage alongside
47% answered-slice error at alpha = 0.10 on a population built to split the two). If the
answered slice itself needs a risk bound, calibrate it directly with
calibrate_selective(), whose guarantee – with probability at least 1 - delta over an
i.i.d. calibration draw, error among answered queries is at most alpha – is the
answered-slice statement, at the cost of quantized thresholds and more calibration data.
Both statements are conditional on their sampling assumption, and both FAIL SILENTLY under
distribution shift: if deployment queries drift from the calibration distribution (new topics,
adversarial inputs, a different upstream data feed), neither the marginal nor the selective
bound holds any longer, and nothing in the serving path detects the drift for you. The
density gate (mixle.task.density) catches inputs the calibration set never resembled;
re-measure escalation_rate and answered-slice agreement on fresh labeled data whenever
the query mix changes.
distill_for_routing (and distill_records_for_routing for structured
records) does steps 2 and 3 in one call – it holds out a calibration slice,
fits the student, and returns an already-decide()-able
CalibratedTaskModel:
from mixle.task import distill_for_routing
calibrated = distill_for_routing(teacher, unlabeled_emails, labels=["spam", "ham"])
For generation rather than classification – e.g. choosing whether to serve
the external checkpoint’s best candidate or escalate –
CalibratedGenerator provides a selective-risk gate. It draws k
candidates, proposes top-score thresholds on one calibration split, and
certifies the accepted-error upper bound on an independent split. Candidate
strings and their count are not a fixed label space, so this surface does not
claim conformal label coverage. Its receipt names the target error,
confidence, accepted certification cases, observed errors, and exact
Bonferroni-corrected binomial upper bound.
Note
CalibratedGenerator (mixle.task.calibrated_generator) is not yet
re-exported from mixle.task’s public surface – import it directly
from its module until it is.
4. Route: Cheap mixle-Native Model First, Expensive External Model Only When Needed¶
Cascade serves traffic with the calibrated local model, escalating to the
external teacher only when the local model is unsure:
from mixle.task import Cascade, CostModel
cascade = Cascade(calibrated, teacher, cost=CostModel(c_frontier=0.03, c_local=0.0001))
label = cascade("free prize inside!!!")
cascade.report() # realized $/request saved vs. teacher-only
cascade.harvested() # (texts, labels) the teacher answered -- feed back into distill()
Router generalizes this to several calibrated tiers (cheapest first,
… -> the external checkpoint as the final, always-answering fallback).
Build it from (name, model, cost) tuples – any tier before the last
just needs a decide(x) method, exactly what calibrated above
already has:
from mixle.task import Router
router = Router([
("local-checkpoint", calibrated, 0.0001), # the calibrated student from step 3, above
("external-checkpoint", teacher, 0.03), # the wrapped external checkpoint, last and most expensive
])
router("free prize inside!!!")
router.report() # per-tier traffic and realized cost
router.harvested() # requests only the external model answered -- re-distill to shrink escalation
A deeper stack (tiny -> small -> … -> external) is the same shape with
more tuples; from_solutions() is a
convenience constructor for the case where each tier is a full
Solution (from solve())
rather than a bare calibrated model – it reads each Solution’s
.cascade.model for you.
In both cases the external, HF/TRL-trained model never needs to change: it is the fallback tier, called exactly the same way it was called in step 1.
This vs. The GradLeaf Pattern¶
There are two, deliberately different, ways an externally-trained model enters mixle:
Pattern |
What happens to the model |
Use when |
|---|---|---|
Bring your own model (this page) |
stays external; called through |
you want to distill a cheap local student from it, calibrate an abstention guarantee, and route around it – without retraining it |
|
loaded into mixle as a real |
you want to keep training the checkpoint itself, inside a mixle model/estimation loop, not just query it |
examples/peft_lora_grad_leaf.py is the concrete receipt for the second pattern: a real
hf-internal-testing/tiny-random-gpt2 checkpoint, LoRA-wrapped with
peft.get_peft_model, dropped into GradLeaf unchanged and fine-tuned
via mixle’s ordinary M-step. That is genuinely different work from this
page’s pattern – no fine-tuning happens here, the checkpoint is called,
not trained.
If you are unsure which you want: if you plan to keep training the external
checkpoint’s weights, use the GradLeaf bridge. If you only want to call
it and cheapen/guarantee/route around those calls, use the callable-teacher
pattern above.
Where Each Receipt Lives¶
Every claim on this page is pinned by an executable test. A skeptical reader should run these directly rather than trust the prose above:
Step |
Claim |
Test file |
|---|---|---|
|
|
|
|
|
|
|
|
|
|
|
|
2+3. One-call routing distill |
|
|
|
|
|
|
|
|
|
one calibrated tier plus teacher fallback; realized cost and harvest |
|
|
N calibrated tiers cheapest-first, teacher as final fallback |
|
GradLeaf contrast |
a real peft-wrapped HF checkpoint fine-tunes unchanged inside
|
|
See Also¶
Task Distillation – the full teacher/student/calibration/escalation model, including numeric, multi-label, and structured teachers.
Task Serving, Routing, and Edge Deployment –
Cascade/Routerin more depth, plus edge deployment and quantization once the student is trained.Neural and LLM Models – the
GradLeafbridge for training a real torch module (including a peft-wrapped HF checkpoint) inside mixle.