| """RubriX: score a response against a rubric's own level descriptions, without text generation. |
| |
| from transformers.dynamic_module_utils import get_class_from_dynamic_module |
| RubriX = get_class_from_dynamic_module("rubrix.RubriX", "CobuildX-AI/RubriX-8B") |
| judge = RubriX.from_pretrained("CobuildX-AI/RubriX-8B") # needs a GPU with ~20 GB (bf16) |
| result = judge.score( |
| criterion="Does the response answer the question correctly and completely?", |
| levels=["Wrong or unrelated.", "Partly correct.", "Correct but incomplete.", "Correct and complete."], |
| response="Paris is the capital of France.", |
| conversation="What is the capital of France?", |
| ) |
| result["probabilities"] # one probability per level, lowest level first |
| result["level"] # index of the most likely level (0-based) |
| result["expected_level"] # probability-weighted level (0-based), useful for ranking |
| |
| How it works: one encoder (Qwen3-8B with this repository's LoRA adapter) embeds the input and each level description |
| separately; two projection heads map them into a shared space; the score of level k is τ · cos(f(input), g(level_k)) |
| and a softmax over the rubric's levels gives the probabilities. Level vectors do not depend on the input, so they are |
| cached: judging many responses on one rubric embeds its levels once. |
| |
| The input text is laid out exactly as in training: |
| |
| [system prompt] (optional) |
| ... |
| [conversation] (optional: the prompt or dialogue the response answers) |
| ... |
| [reference answer] (optional) |
| ... |
| [response to evaluate] |
| ... |
| |
| <criterion> |
| |
| Requirements: torch, transformers, peft, safetensors, huggingface_hub. |
| """ |
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| BASE_MODEL = "Qwen/Qwen3-8B" |
| BASE_REVISION = "b968826d9c46dd6066d109eabc6255188de91218" |
|
|
|
|
| class _Head(nn.Module): |
| """hidden -> width -> width -> proj MLP with GELU and LayerNorm (the CLM head layout).""" |
|
|
| def __init__(self, hidden: int, width: int, depth: int, proj: int, layernorm: bool): |
| super().__init__() |
| self.inp = nn.Linear(hidden, width) |
| self.hidden = nn.ModuleList(nn.Linear(width, width) for _ in range(depth - 2)) |
| self.norms = nn.ModuleList((nn.LayerNorm(width) if layernorm else nn.Identity()) for _ in range(depth - 2)) |
| self.out = nn.Linear(width, proj) |
| self.act = nn.GELU() |
|
|
| def forward(self, x): |
| x = self.act(self.inp(x)) |
| for lin, nrm in zip(self.hidden, self.norms): |
| x = self.act(nrm(lin(x))) |
| return self.out(x) |
|
|
|
|
| def build_input(response: str, criterion: str, conversation: str | None = None, system_prompt: str | None = None, |
| reference_answer: str | None = None) -> str: |
| """The text the input encoder sees: labelled sections, then the criterion after a blank line.""" |
| parts = {"system prompt": system_prompt, "conversation": conversation, "reference answer": reference_answer, |
| "response to evaluate": response} |
| state = "\n\n".join(f"[{k}]\n{v}" for k, v in parts.items() if isinstance(v, str) and v.strip()) |
| s, c = state.strip(), (criterion or "").strip() |
| return f"{s}\n\n{c}" if s and c else (s or c) |
|
|
|
|
| class RubriX: |
| def __init__(self, tokenizer, encoder, state_head, level_head, scale: float, device, max_tokens: int = 32768): |
| self.tok, self.encoder = tokenizer, encoder |
| self.state_head, self.level_head = state_head, level_head |
| self.scale, self.device, self.max_tokens = scale, device, max_tokens |
| self._level_cache: dict[str, torch.Tensor] = {} |
|
|
| @classmethod |
| def from_pretrained(cls, repo_id: str = "CobuildX-AI/RubriX-8B", device: str | None = None, |
| dtype: torch.dtype = torch.bfloat16, merge: bool = True, token: str | bool | None = None, |
| max_tokens: int | None = None): |
| """Load Qwen3-8B (pinned revision), apply the LoRA adapter and load the heads. `repo_id` may be a local folder. |
| `max_tokens` caps each input (default 32,768, Qwen3-8B's native context; RubriX was trained on up to 8,192).""" |
| from huggingface_hub import snapshot_download |
| from peft import PeftModel |
| from safetensors.torch import load_file |
| from transformers import AutoModel, AutoTokenizer |
|
|
| path = Path(repo_id) if Path(repo_id).is_dir() else Path(snapshot_download(repo_id, token=token)) |
| cfg = json.loads((path / "rubrix_config.json").read_text()) |
| device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| tok = AutoTokenizer.from_pretrained(cfg["base_model"], revision=cfg["base_revision"]) |
| tok.padding_side = "right" |
| base = AutoModel.from_pretrained(cfg["base_model"], revision=cfg["base_revision"], dtype=dtype) |
| enc = PeftModel.from_pretrained(base, str(path / "adapter")) |
| enc = enc.merge_and_unload() if merge else enc |
| enc = enc.to(device).eval() |
| h = cfg["heads"] |
| heads = load_file(str(path / "heads.safetensors")) |
| state_head = _Head(h["hidden_size"], h["width"], h["depth"], h["projection_dim"], h["layernorm"]) |
| level_head = _Head(h["hidden_size"], h["width"], h["depth"], h["projection_dim"], h["layernorm"]) |
| state_head.load_state_dict({k[len("state_head."):]: v for k, v in heads.items() if k.startswith("state_head.")}) |
| level_head.load_state_dict({k[len("level_head."):]: v for k, v in heads.items() if k.startswith("level_head.")}) |
| scale = min(float(torch.exp(heads["logit_scale"].float())), h.get("max_scale", 100.0)) |
| return cls(tok, enc, state_head.to(device).float().eval(), level_head.to(device).float().eval(), scale, device, |
| max_tokens or cfg.get("max_tokens", 32768)) |
|
|
| @torch.no_grad() |
| def embed(self, texts: list[str], batch_size: int = 8) -> torch.Tensor: |
| """L2-normalised final hidden state of each text's last token.""" |
| out = [] |
| for i in range(0, len(texts), batch_size): |
| b = self.tok(texts[i:i + batch_size], return_tensors="pt", padding=True, truncation=True, |
| max_length=self.max_tokens).to(self.device) |
| h = self.encoder(**b, use_cache=False).last_hidden_state |
| last = b["attention_mask"].sum(1) - 1 |
| out.append(F.normalize(h[torch.arange(h.shape[0], device=h.device), last].float(), dim=-1)) |
| return torch.cat(out) |
|
|
| @torch.no_grad() |
| def _levels(self, levels: list[str]) -> torch.Tensor: |
| new = [t for t in dict.fromkeys(levels) if t not in self._level_cache] |
| if new: |
| for t, v in zip(new, F.normalize(self.level_head(self.embed(new)), dim=-1)): |
| self._level_cache[t] = v |
| return torch.stack([self._level_cache[t] for t in levels]) |
|
|
| def clear_cache(self): |
| self._level_cache.clear() |
|
|
| @torch.no_grad() |
| def score_batch(self, items: list[dict], batch_size: int = 8) -> list[dict]: |
| """items: dicts with criterion, levels (lowest first), response, and optional conversation, system_prompt, |
| reference_answer. Returns one result per item.""" |
| texts = [build_input(it["response"], it["criterion"], it.get("conversation"), it.get("system_prompt"), |
| it.get("reference_answer")) for it in items] |
| z = F.normalize(self.state_head(self.embed(texts, batch_size)), dim=-1) |
| results = [] |
| for zi, it in zip(z, items): |
| if len(it["levels"]) < 2: |
| raise ValueError("a rubric needs at least two levels") |
| logits = self.scale * (self._levels(list(it["levels"])) @ zi) |
| p = torch.softmax(logits, dim=-1) |
| results.append({"probabilities": p.tolist(), "level": int(p.argmax()), |
| "expected_level": float((p * torch.arange(len(p), device=p.device)).sum())}) |
| return results |
|
|
| def score(self, criterion: str, levels: list[str], response: str, conversation: str | None = None, |
| system_prompt: str | None = None, reference_answer: str | None = None) -> dict: |
| """Score one response. `levels` are the rubric's level descriptions, lowest first.""" |
| return self.score_batch([{"criterion": criterion, "levels": levels, "response": response, |
| "conversation": conversation, "system_prompt": system_prompt, |
| "reference_answer": reference_answer}])[0] |
|
|
