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 complete()/next_logprobs() as a black-box teacher

you want to distill a cheap local student from it, calibrate an abstention guarantee, and route around it – without retraining it

GradLeaf bridge (see Neural and LLM Models and examples/peft_lora_grad_leaf.py)

loaded into mixle as a real torch.nn.Module and fine-tuned in-process (e.g. a peft-wrapped HF checkpoint’s LoRA adapters trained via mixle.inference.estimation.optimize)

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

  1. Callable teacher

CallableLLM/OpenAICompatLLM wrap an arbitrary function/endpoint as an LLM; llm_labeler constrains it to a label set

mixle/tests/task_llm_test.py

  1. VLM sibling

CallableVLM/OpenAICompatVLM wrap an image-conditioned scorer

mixle/tests/task_vlm_test.py

  1. Distill

distill/distill_from_labels reproduce teacher labels in a small local student; agreement is recorded

mixle/tests/task_distill_test.py, mixle/tests/distill_methods_test.py

  1. Active distill

active_distill spends a fixed label budget via acquisition

mixle/tests/task_active_test.py, mixle/tests/doe_active_test.py

2+3. One-call routing distill

distill_for_routing/distill_records_for_routing return an already-calibrated model

mixle/tests/task_distill_routing_test.py

  1. Calibrate

CalibratedTaskModel gives 1 - alpha conformal coverage and an honest ESCALATE on ambiguous/OOD input

mixle/tests/task_calibrate_test.py

  1. Calibrated generation

CalibratedGenerator gives a held-out selective-risk certificate for serving the top sampled candidate or abstaining

mixle/tests/task_calibrated_generator_test.py

  1. Cascade

one calibrated tier plus teacher fallback; realized cost and harvest

mixle/tests/task_cascade_test.py

  1. Router

N calibrated tiers cheapest-first, teacher as final fallback

mixle/tests/task_router_test.py

GradLeaf contrast

a real peft-wrapped HF checkpoint fine-tunes unchanged inside GradLeaf

mixle/tests/grad_control_test.py (AdapterThroughTheBridgeTest), examples/peft_lora_grad_leaf.py

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/Router in more depth, plus edge deployment and quantization once the student is trained.

  • Neural and LLM Models – the GradLeaf bridge for training a real torch module (including a peft-wrapped HF checkpoint) inside mixle.