import refrom transformers import AutoModelForCausalLM, AutoTokenizer REPO = "opsd-genrm/SDRM-Qwen3-4B-mask70"tok = AutoTokenizer.from_pretrained(REPO)model = AutoModelForCausalLM.from_pretrained( REPO, torch_dtype="bfloat16", device_map="auto") PROMPT = """You are an impartial judge tasked with determining which of two assistant responses is better for the given context. Below is a context (a user query or a conversation between the user and an assistant) and two assistant responses to that context. [Start of Context]{context}[End of Context] [Start of Assistant A's Response]{response_a}[End of Assistant A's Response] [Start of Assistant B's Response]{response_b}[End of Assistant B's Response] Identify the quality dimensions that matter most for this specific task, then evaluate and compare the two assistant responses step by step across those dimensions. When correctness matters, solve the problem yourself and check each response for any errors. After your analysis, determine which response is better overall and provide your final verdict (A or B only) in <verdict>...</verdict>.""" def format_context(messages: list[dict]) -> str: """Wrap a (multi-turn) dialog as alternating <user>/<assistant> blocks. `messages` is a list of {"role": "user"|"assistant", "content": str}, must alternate starting with "user", and must end on a "user" turn. """ parts = [] for i, m in enumerate(messages): expected = "user" if i % 2 == 0 else "assistant" assert m["role"] == expected, "roles must alternate user/assistant/..." parts.append(f"<{m['role']}>\n{m['content'].strip()}\n</{m['role']}>") return "\n\n".join(parts) def format_response(text: str) -> str: return f"<assistant>\n{text.strip()}\n</assistant>" def parse_verdict(text: str) -> str | None: """Prefers the last <verdict>...</verdict> block: returns the last 'A' or 'B' word inside it. If the verdict tag is missing or malformed, falls back to the last standalone 'A'/'B' anywhere in the generated text. Returns None only when no A/B token appears at all. """ blocks = re.findall(r"<verdict>(.*?)</verdict>", text, re.DOTALL) if blocks: ab = re.findall(r"\b(A|B)\b", blocks[-1]) if ab: return ab[-1] ab = re.findall(r"\b(A|B)\b", text) return ab[-1] if ab else None def judge(context_messages: list[dict], response_a: str, response_b: str) -> str | None: user_msg = PROMPT.format( context=format_context(context_messages), response_a=format_response(response_a), response_b=format_response(response_b), ) inputs = tok.apply_chat_template( [{"role": "user", "content": user_msg}], add_generation_prompt=True, return_tensors="pt", ).to(model.device) out = model.generate( inputs, max_new_tokens=8192, do_sample=True, temperature=0.7, top_p=0.8, top_k=20, ) text = tok.decode(out[0, inputs.shape[-1]:], skip_special_tokens=True) return parse_verdict(text) # Single-turn:verdict = judge( context_messages=[{"role": "user", "content": "What is the capital of France?"}], response_a="The capital of France is Paris.", response_b="Lyon.",)print(verdict) # -> "A"