File size: 3,305 Bytes
ff2b400 | 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 | """Minimal EditRank-Multidomain-4B usage: score candidate rewrites of a document for a query and pick the top-1.
The score is logit("yes") - logit("no") at the last position, exactly as in evaluation.
Higher = the model expects the revised document to rank higher in an answer engine's source list for the query.
The unchanged original is a valid candidate: include it, and keep it when it scores highest.
Research Preview: do NOT auto-publish the top-1 edit. Every candidate needs independent factual validation.
"""
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL_ID = "ThakiCloud/EditRank-Multidomain-4B"
PREFIX = ('<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. '
'Note that the answer can only be "yes" or "no".<|im_end|>\n<|im_start|>user\n')
SUFFIX = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
EDIT_INSTR = ("Given a search query, an original document and a revised version of it, judge whether the revised version would be "
"selected and ranked higher by an AI answer engine for this query than the original, while keeping every fact of "
"the original unchanged.")
def clip(s, n=1500):
return s if len(s) <= n else s[:n] + " ..."
def build(query, original, revised):
doc = f"ORIGINAL (excerpt): {clip(original, 600)}\n\nREVISED: {clip(revised)}"
return f"{PREFIX}<Instruct>: {EDIT_INSTR}\n<Query>: {query}\n<Document>: {doc}{SUFFIX}"
@torch.no_grad()
def score(model, tok, query, original, candidates, max_len=1280):
yes_id, no_id = tok.convert_tokens_to_ids("yes"), tok.convert_tokens_to_ids("no")
enc = tok([build(query, original, c) for c in candidates], padding=True, truncation=True,
max_length=max_len, return_tensors="pt").to(model.device)
logits = model(**enc).logits[:, -1, :]
return (logits[:, yes_id].float() - logits[:, no_id].float()).tolist()
if __name__ == "__main__":
tok = AutoTokenizer.from_pretrained(MODEL_ID, padding_side="left")
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16,
device_map="cuda" if torch.cuda.is_available() else None).eval()
query = "how do I limit the memory a container can use in a kubernetes pod"
original = ("Resource limits. Each container can declare resources. Set resources.limits.memory in the container spec. "
"If the container exceeds the limit it may be terminated.")
candidates = [
original,
"To limit a container's memory in a Pod, set resources.limits.memory in that container's spec. "
"If the container exceeds its memory limit, it may be terminated.",
"To limit a container's memory in a Pod, set resources.limits.memory. Kubernetes 1.31 enforces limits "
"within 50 ms, so 99.9% of clusters never see out-of-memory kills.", # invents facts
]
s = score(model, tok, query, original, candidates)
for c, v in sorted(zip(candidates, s), key=lambda x: -x[1]):
print(f"{v:+.3f} {c[:90]}")
print("The third candidate invents a version, a latency and a statistic; the model is not a fact checker. "
"Validate every fact before using any edit.")
|