#!/usr/bin/env python3 """ Official LongMemEval grading — replicates xiaowu0162/LongMemEval/src/evaluation/evaluate_qa.py exactly. Uses GPT-4o as the grader model, one call per question, same templates. """ import argparse import json import os import sys import backoff import openai from openai import OpenAI import numpy as np def get_anscheck_prompt(task, question, answer, response, abstention=False): if not abstention: if task in ['single-session-user', 'single-session-assistant', 'multi-session']: template = "I will give you a question, a correct answer, and a response from a model. Please answer yes if the response contains the correct answer. Otherwise, answer no. If the response is equivalent to the correct answer or contains all the intermediate steps to get the correct answer, you should also answer yes. If the response only contains a subset of the information required by the answer, answer no. \n\nQuestion: {}\n\nCorrect Answer: {}\n\nModel Response: {}\n\nIs the model response correct? Answer yes or no only." return template.format(question, answer, response) elif task == 'temporal-reasoning': template = "I will give you a question, a correct answer, and a response from a model. Please answer yes if the response contains the correct answer. Otherwise, answer no. If the response is equivalent to the correct answer or contains all the intermediate steps to get the correct answer, you should also answer yes. If the response only contains a subset of the information required by the answer, answer no. In addition, do not penalize off-by-one errors for the number of days. If the question asks for the number of days/weeks/months, etc., and the model makes off-by-one errors (e.g., predicting 19 days when the answer is 18), the model's response is still correct. \n\nQuestion: {}\n\nCorrect Answer: {}\n\nModel Response: {}\n\nIs the model response correct? Answer yes or no only." return template.format(question, answer, response) elif task == 'knowledge-update': template = "I will give you a question, a correct answer, and a response from a model. Please answer yes if the response contains the correct answer. Otherwise, answer no. If the response contains some previous information along with an updated answer, the response should be considered as correct as long as the updated answer is the required answer.\n\nQuestion: {}\n\nCorrect Answer: {}\n\nModel Response: {}\n\nIs the model response correct? Answer yes or no only." return template.format(question, answer, response) elif task == 'single-session-preference': template = "I will give you a question, a rubric for desired personalized response, and a response from a model. Please answer yes if the response satisfies the desired response. Otherwise, answer no. The model does not need to reflect all the points in the rubric. The response is correct as long as it recalls and utilizes the user's personal information correctly.\n\nQuestion: {}\n\nRubric: {}\n\nModel Response: {}\n\nIs the model response correct? Answer yes or no only." return template.format(question, answer, response) else: raise NotImplementedError(f"Unknown task: {task}") else: template = "I will give you an unanswerable question, an explanation, and a response from a model. Please answer yes if the model correctly identifies the question as unanswerable. The model could say that the information is incomplete, or some other information is given but the asked information is not.\n\nQuestion: {}\n\nExplanation: {}\n\nModel Response: {}\n\nDoes the model correctly identify the question as unanswerable? Answer yes or no only." return template.format(question, answer, response) @backoff.on_exception(backoff.expo, (openai.RateLimitError, openai.APIError)) def grade_one(client, model, prompt): completion = client.chat.completions.create( model=model, messages=[{"role": "user", "content": prompt}], n=1, temperature=0, max_tokens=10, ) return completion.choices[0].message.content.strip() def main(): p = argparse.ArgumentParser() p.add_argument("--input", required=True, help="Replay JSONL (must have question_id, question_type, question, ground_truth, c137_answer)") p.add_argument("--output", required=True, help="Output JSONL with grades") p.add_argument("--model", default="gpt-4o", help="Grader model (default: gpt-4o)") args = p.parse_args() api_key = os.getenv("OPENAI_API_KEY") if not api_key: print("OPENAI_API_KEY not set", file=sys.stderr) sys.exit(1) client = OpenAI(api_key=api_key) entries = [] for line in open(args.input): line = line.strip() if line: entries.append(json.loads(line)) print(f"Loaded {len(entries)} entries") qtypes = set() qtype2acc = {} results = [] with open(args.output, "w") as out: for i, e in enumerate(entries): qid = e["question_id"] qtype = e["question_type"] qtypes.add(qtype) if qtype not in qtype2acc: qtype2acc[qtype] = [] question = e["question"] answer = e["ground_truth"] response = e.get("c137_answer", "") abstention = "_abs" in qid prompt = get_anscheck_prompt(qtype, question, answer, response, abstention=abstention) eval_response = grade_one(client, args.model, prompt) label = "yes" in eval_response.lower() result = { "question_id": qid, "question_type": qtype, "label": label, "raw": eval_response, } out.write(json.dumps(result) + "\n") out.flush() qtype2acc[qtype].append(1 if label else 0) results.append(result) if (i + 1) % 50 == 0 or (i + 1) == len(entries): total_acc = np.mean([r["label"] for r in results]) print(f" {i+1}/{len(entries)} ({100*total_acc:.1f}%)", flush=True) # Final report total = sum(len(v) for v in qtype2acc.values()) correct = sum(sum(v) for v in qtype2acc.values()) print(f"\nTotal: {correct}/{total} = {100*correct/total:.1f}%") for k in sorted(qtype2acc.keys()): v = qtype2acc[k] print(f" {k:30s} {100*np.mean(v):.1f}% ({sum(v)}/{len(v)})") if __name__ == "__main__": main()