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