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]
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 ..."))