qwen3-embedding-8b-ielts-band-lora

IELTS Overall-Band Scorer — Qwen3-Embedding-8B + LoRA

A LoRA adapter on Qwen/Qwen3-Embedding-8B plus a small heteroscedastic regression head that predicts the overall IELTS Writing band (0–9, half-steps) of an essay, with a calibrated uncertainty.

⚠️ This is not a generative / chat model. The encoder produces an embedding; a tiny regression head (shipped here as head_np.npz) turns that embedding into a band. Use it as an embedding model, then run the head.

Results (held-out, 49,071 essays)

Strongest in the dense middle-to-upper region (bands 4.5–5.0, 7.0–8.0: MAE 0.43–0.54). Weakest at the extremes (band 9: 0.87) — the residual there is dominated by label noise in the training source (independent expert graders themselves disagree ~0.4 bands), not model capacity.

How it works

essay ─▶ Qwen3-Embedding-8B + this LoRA ─▶ 4096-d last-token embedding ─▶ head ─▶ (μ, σ) ─▶ band
        (last-token pool, NO L2 normalization)

• LoRA: r=16, α=32, on q/k/v/o + gate/up/down projections.
• Head (head_np.npz): Linear(4096→512) → GELU → Linear(512→2) → (μ, log σ²). Band = round-to-0.5(μ); σ = exp(½·logσ²)·T with temperature T=1.076 (see meta.json).
• Critical: the head is trained on the raw, un-normalized last-token hidden state. When serving the embedding you must use last-token pooling with normalization off.

Usage

Quick start (transformers + peft)

import numpy as np, torch, torch.nn as nn, math, json
from transformers import AutoModel, AutoTokenizer
from peft import PeftModel
from huggingface_hub import hf_hub_download

REPO = "ndtran0101/qwen3-embedding-8b-ielts-band-lora"
tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-Embedding-8B"); tok.padding_side = "left"
base = AutoModel.from_pretrained("Qwen/Qwen3-Embedding-8B", torch_dtype=torch.bfloat16)
model = PeftModel.from_pretrained(base, REPO).eval().cuda()

h = np.load(hf_hub_download(REPO, "head_np.npz")); meta = json.load(open(hf_hub_download(REPO, "meta.json")))
def gelu(x): return 0.5*x*(1+np.vectorize(math.erf)(x/math.sqrt(2)))

def score(essay, question=""):
    q = question.strip(); body = essay.strip()
    text = f"Instruct: {meta['instruct']}\nQuery: {q}\n\n{body}" + tok.eos_token
    enc = tok([text], return_tensors="pt", truncation=True, max_length=768).to("cuda")
    with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
        emb = model(**enc).last_hidden_state[:, -1, :].float().cpu().numpy()[0]  # raw, un-normalized
    o = gelu(emb @ h["0_weight"].T + h["0_bias"]) @ h["3_weight"].T + h["3_bias"]
    mu = float(o[0]); sigma = math.exp(0.5*float(np.clip(o[1], -6, 4)))*meta["temperature"]
    return {"band": round(min(max(mu,0),9)*2)/2, "mu": round(mu,3), "sigma": round(sigma,3)}

print(score("Living in a city has many advantages ...", question="Some people prefer city life ..."))

Serving with vLLM

Serve the LoRA-adapted encoder as an embedding endpoint (last-token pool, no normalize), then run the head client-side with the bundled score_band.py:

docker run --restart unless-stopped -d --name qwen-embeddings-8b \
  --gpus '"device=0"' -v ~/.cache:/root/.cache -p 8000:8000 --ipc host \
  vllm/vllm-openai:nightly \
  Qwen/Qwen3-Embedding-8B \
  --served-model-name qwen/qwen3-embedding-8b \
  --runner pooling \
  --pooler-config '{"pooling_type":"LAST","use_activation":false}' \
  --enable-lora \
  --lora-modules qwen/qwen3-embedding-8b:ielts-band=ndtran0101/qwen3-embedding-8b-ielts-band-lora \
  --max-lora-rank 16

Then embed with model="qwen/qwen3-embedding-8b:ielts-band" and feed the vector to the head. See score_band.py.

Files

Training

Fine-tuned end-to-end (LoRA on the encoder + head) on ndtran0101/writing9-ielts-essays (114k train). Objective: Gaussian β-NLL (heteroscedastic mean–variance regression; Nix & Weigend 1994, Kendall & Gal 2017) with inverse-frequency class weighting and post-hoc temperature calibration of σ. Config: lr 2e-4, class-weight power 1.0, β 0.5, last-token pooling, 2 epochs, plus 2,500 LLM-polished band-9 examples for tail augmentation.

Limitations

Trained against crowd-sourced band labels that are noisy at the extremes; treat single-band precision for bands ≤4.5 and ≥8.5 with caution and prefer the ±σ range. English only.