File size: 8,718 Bytes
248975b c07a846 248975b c07a846 248975b c07a846 248975b c07a846 248975b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | """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" # the exact revision RubriX was trained on
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" # pooling reads the last real token of each row
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]
|