qwen-swan-sig-2b

Fine-tuned Qwen3.5-2B lightweight variant of Nbeau/qwen-swan-sig-4b, for the same Swan / SCADE / Lustre signature-prediction task. Use this model when inference cost or on-device deployment matters (~4 GB bf16, ~1.5 GB at INT4).

Task

Given a target operator name (+ optional NL description) and a few neighboring operators of the same module, output a JSON object {"inputs": [...], "outputs": [...]} where each item is {"name", "type"}.

Quick use

Same API and prompt format as the 4B model - just swap the repo name:

import json, torch
from transformers import AutoModelForCausalLM, AutoTokenizer

REPO = "Nbeau/qwen-swan-sig-2b"
SYSTEM_PROMPT = (
    "You are a Swan/Lustre/SCADE language expert. "
    "Given context operators from a module and a target operator name, "
    "predict its inputs and outputs as a JSON object with keys: inputs, outputs. "
    "Each input/output is {name, type}. "
    "Respond with only the JSON object, no explanation.\n\n"
    "Swan type-system reference:\n"
    "- Primitive types: bool, int8, int16, int32, int64, float32, float64, char.\n"
    "- Array types: T^N means array of type T with size N. Multi-dimensional: T^X^Y.\n"
    "- Generic type parameters: 'T, 'U, 'F (polymorphic).\n"
    "- Generic constants: <<PARAM1, PARAM2>> after the operator name."
)

USER_PROMPT = """[CONTEXT]
<module>libraries_Float</module>
function LE (a: 'T; b: 'T; e: 'T;) returns (comp: bool;)
function GE (a: 'T; b: 'T; e: 'T;) returns (comp: bool;)
function LT (a: 'T; b: 'T; e: 'T;) returns (comp: bool;)
[PREDICT]
Name: GT
Description: Strictly greater-than with tolerance epsilon."""

tok = AutoTokenizer.from_pretrained(REPO)
model = AutoModelForCausalLM.from_pretrained(REPO, torch_dtype=torch.bfloat16, device_map="auto")

text = tok.apply_chat_template(
    [{"role": "system", "content": SYSTEM_PROMPT},
     {"role": "user",   "content": USER_PROMPT}],
    tokenize=False, enable_thinking=False, add_generation_prompt=True,
)
inputs = tok(text, return_tensors="pt").to(model.device)

out = model.generate(
    **inputs,
    max_new_tokens=384,
    do_sample=False,
    repetition_penalty=1.05,
    no_repeat_ngram_size=20,
    pad_token_id=tok.eos_token_id,
)
print(json.loads(tok.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()))
# {"inputs": [{"name": "a", "type": "'T"}, {"name": "b", "type": "'T"}, {"name": "e", "type": "'T"}],
#  "outputs": [{"name": "comp", "type": "bool"}]}

See the 4B card for the full prompt format reference.

Results vs 4B sibling

Held-out test: 39 Swan operators from 4 entire modules never seen during training.

Metric	2B (this)	4B sibling	Delta
exact_match (overall)	0.231	0.282	-0.05
structural_f1 (overall)	0.406	0.445	-0.04
type_f1 (overall)	0.628	0.649	-0.02
shape_accuracy (overall)	0.487	0.410	+0.08
sparse exact_match	0.100	0.300	-0.20
metier structural_f1	0.219	0.260	-0.04
EM count	9 / 39	11 / 39	-2
Runaway	0	0	=

The 2B is comparable on full-context prompts but markedly weaker in the sparse regime (1-3 neighbors). Prefer the 4B when you can afford the extra compute.

Training

Same data and method as the 4B sibling, with these differences:

• Method. LoRA r=16, alpha=32, dropout=0.10 (~12 M trainable params), then merged.
• Hyperparameters. LR=5e-5 (the 2B does not benefit from the 4B's LR=1e-4 - verified by ablation: at 1e-4 the 2B regressed on every metier metric), cosine, warmup 200, effective batch 32 (bs=4 x grad_accum=8), seq_len=2048, weight_decay=0.05, eval/save every 50 steps, early-stop patience 8.
• Best-checkpoint selection. Same custom structural_f1 callback as the 4B.
• Best ckpt. Step 250 / epoch 0.56, eval_structural_f1 = 0.329.
• Hardware. 1x NVIDIA H100 NVL 96 GB, ~5 h.

Limitations

All limitations of the 4B model apply, plus:

• Sparse-regime quality is ~3x lower than the 4B (10 % vs 30 % EM). If you have neighbors but the task feels hard, route to the 4B.
• Slightly less robust on rare metier types (AStarStats_t, T_Mode_Level).

License

Apache 2.0 (inherited from Qwen/Qwen3.5-2B).