| |
| """Generate model outputs for a VERL SFT validation parquet file.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import re |
| from pathlib import Path |
|
|
| import pandas as pd |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| THINK_BLOCK_RE = re.compile(r"(?is)<think>(.*?)</think>") |
| ANSWER_BLOCK_RE = re.compile(r"(?is)<answer>(.*?)</answer>") |
|
|
|
|
| def split_tagged_response(text: str) -> dict[str, object]: |
| think_match = THINK_BLOCK_RE.search(text) |
| answer_match = ANSWER_BLOCK_RE.search(text) |
| return { |
| "think": think_match.group(1).strip() if think_match else "", |
| "answer": answer_match.group(1).strip() if answer_match else "", |
| "has_think": think_match is not None, |
| "has_answer": answer_match is not None, |
| } |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model", required=True, help="HF model directory or model id.") |
| parser.add_argument("--adapter", help="Optional PEFT/LoRA adapter directory.") |
| parser.add_argument("--tokenizer", help="Optional tokenizer path. Defaults to --model.") |
| parser.add_argument("--data", default="data/hep_sft/val.parquet") |
| parser.add_argument("--output", default="data/hep_sft/validation_outputs.jsonl") |
| parser.add_argument("--limit", type=int, default=32) |
| parser.add_argument("--max-new-tokens", type=int, default=512) |
| args = parser.parse_args() |
|
|
| model_path = Path(args.model) |
| data_path = Path(args.data) |
| output_path = Path(args.output) |
| output_path.parent.mkdir(parents=True, exist_ok=True) |
|
|
| tokenizer = AutoTokenizer.from_pretrained(args.tokenizer or model_path, trust_remote_code=True) |
| model = AutoModelForCausalLM.from_pretrained( |
| model_path, |
| torch_dtype=torch.bfloat16, |
| device_map="auto", |
| trust_remote_code=True, |
| ) |
| if args.adapter: |
| from peft import PeftModel |
|
|
| model = PeftModel.from_pretrained(model, args.adapter) |
| model.eval() |
|
|
| df = pd.read_parquet(data_path) |
| if args.limit > 0: |
| df = df.head(args.limit) |
|
|
| with output_path.open("w") as out: |
| for row in df.to_dict("records"): |
| messages = row["messages"] |
| prompt_messages = [m for m in messages if m["role"] != "assistant"] |
| reference = next((m["content"] for m in messages if m["role"] == "assistant"), "") |
|
|
| inputs = tokenizer.apply_chat_template( |
| prompt_messages, |
| add_generation_prompt=True, |
| tokenize=True, |
| return_tensors="pt", |
| ).to(model.device) |
|
|
| with torch.no_grad(): |
| generated = model.generate( |
| inputs, |
| max_new_tokens=args.max_new_tokens, |
| do_sample=False, |
| temperature=None, |
| top_p=None, |
| pad_token_id=tokenizer.eos_token_id, |
| ) |
|
|
| output_ids = generated[0, inputs.shape[-1] :] |
| prediction = tokenizer.decode(output_ids, skip_special_tokens=True).strip() |
| prediction_parts = split_tagged_response(prediction) |
| reference_parts = split_tagged_response(reference) |
| record = { |
| "id": row.get("id"), |
| "arxiv_id": row.get("arxiv_id"), |
| "task_type": row.get("task_type"), |
| "target_type": row.get("target_type"), |
| "target_category": row.get("target_category"), |
| "target_process_id": row.get("target_process_id"), |
| "target_background": row.get("target_background"), |
| "prompt": prompt_messages[-1]["content"], |
| "prediction": prediction, |
| "reference": reference, |
| "prediction_think": prediction_parts["think"], |
| "prediction_answer": prediction_parts["answer"], |
| "prediction_has_think": prediction_parts["has_think"], |
| "prediction_has_answer": prediction_parts["has_answer"], |
| "reference_think": reference_parts["think"], |
| "reference_answer": reference_parts["answer"], |
| "reference_has_think": reference_parts["has_think"], |
| "reference_has_answer": reference_parts["has_answer"], |
| } |
| out.write(json.dumps(record, ensure_ascii=False) + "\n") |
| print(f"wrote {record['id']}") |
|
|
| print(f"Wrote {output_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
