MOF-deprecated / script /infer_worker.py

Upload folder using huggingface_hub

24c3614 verified 9 months ago

15.5 kB

import os
import json
import math
import time
import csv
import logging
import argparse
from typing import List, Dict, Any, Tuple, Optional

import torch
import torch.nn as nn
from tqdm import tqdm

from transformers import AutoTokenizer, AutoModel
from peft import LoraConfig, get_peft_model

# Flash SDP (Optional)
torch.backends.cuda.enable_flash_sdp(True)
torch.backends.cuda.enable_mem_efficient_sdp(True)
torch.backends.cuda.enable_math_sdp(False)

# =========================
# logging
# =========================
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
logger = logging.getLogger("infer_worker")

# =========================
# System Prompt
# =========================
SYSTEM_PROMPT = """
Act as an expert in reticular chemistry. You will receive reaction conditions as a JSON object with the fields: metal_precursor, 
organic_linker, modulator, solvent, metal_concentration_mM, M_L_ratio, temperature_C, and time_h. Based on these inputs, output exactly one uppercase label:
'P' if the conditions are likely to yield a crystalline metal-organic framework under experimental conditions, or 'N' if not..
"""

# =========================
# Data Loading
# =========================
def load_messages_from_jsonl(path: str) -> List[List[Dict[str, Any]]]:
    all_messages = []
    with open(path, "r", encoding="utf-8-sig") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            data = json.loads(line)
            if "messages" in data:
                all_messages.append(data["messages"])
    return all_messages

def standardize_json_input(json_string: str) -> str:
    try:
        data = json.loads(json_string)
        keep = {
            "metal_precursor": data.get("metal_precursor"),
            "organic_linker": data.get("organic_linker"),
            "modulator": data.get("modulator"),
            "solvent": data.get("solvent"),
            "metal_concentration_mM": data.get("metal_concentration_mM"),
            "M_L_ratio": data.get("M_L_ratio"),
            "temperature_C": data.get("temperature_C"),
            "time_h": data.get("time_h"),
        }
        return json.dumps(keep, ensure_ascii=False)
    except Exception:
        return json_string

def char_norm(s: Optional[str]) -> Optional[str]:
    s = (s or "").strip().upper()
    return s if s in ("P", "N") else None

def standardize_messages(messages: List[Dict[str, Any]]) -> Tuple[List[Dict[str, str]], Optional[str]]:
    user_content = None
    assistant_content = None
    for m in messages:
        if m.get("role") == "user":
            user_content = m.get("content")
        elif m.get("role") == "assistant":
            assistant_content = m.get("content")

    user_content = standardize_json_input(user_content or "")
    standardized = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_content},
    ]
    gold = char_norm(assistant_content)
    return standardized, gold

def get_user_before_last_assistant(messages):
    if not messages:
        return None, None
    last_asst_idx = None
    for i in range(len(messages) - 1, -1, -1):
        if messages[i].get("role") == "assistant":
            last_asst_idx = i
            break
    if last_asst_idx is None:
        return None, None
    gold_char = char_norm(messages[last_asst_idx].get("content"))
    user_raw = None
    for j in range(last_asst_idx - 1, -1, -1):
        if messages[j].get("role") == "user":
            user_raw = messages[j].get("content")
            break
    return user_raw, gold_char

# =========================
# Model
# =========================
def build_model_with_lora(base_model_name: str, cache_dir: str, device: torch.device):
    lora_config = LoraConfig(
        r=8,
        lora_alpha=16,
        target_modules=["q_proj", "v_proj", "o_proj"],
        lora_dropout=0.01,
        bias="none",
        task_type="SEQ_CLS"
    )

    backbone = AutoModel.from_pretrained(
        base_model_name,
        torch_dtype=torch.bfloat16,
        device_map=None,  # 强制单卡
        cache_dir=cache_dir,
    ).to(device)

    backbone = get_peft_model(backbone, lora_config)

    class QwenWithLoRAForBinaryClassification(nn.Module):
        def __init__(self, backbone):
            super().__init__()
            self.backbone = backbone
            self.config = backbone.config
            hidden_size = backbone.config.hidden_size
            self.classifier = nn.Linear(hidden_size, 2)

        def forward(self, input_ids, attention_mask=None, labels=None):
            outputs = self.backbone(
                input_ids=input_ids,
                attention_mask=attention_mask,
                return_dict=True,
            )
            mask = attention_mask.unsqueeze(-1)
            pooled = (outputs.last_hidden_state * mask).sum(1) / mask.sum(1)

            if (self.classifier.weight.device != pooled.device) or (self.classifier.weight.dtype != pooled.dtype):
                self.classifier = self.classifier.to(device=pooled.device, dtype=pooled.dtype)

            logits = self.classifier(pooled)
            loss = None
            if labels is not None:
                loss = nn.CrossEntropyLoss()(logits, labels)
            return {"loss": loss, "logits": logits}

    return QwenWithLoRAForBinaryClassification(backbone).to(device)

def load_trained_weights(model: nn.Module, ckpt_dir: str):
    candidates = [
        os.path.join(ckpt_dir, "pytorch_model.bin"),
        os.path.join(ckpt_dir, "model.safetensors"),
        os.path.join(ckpt_dir, "adapter_model.bin"),
    ]
    weight_path = None
    for p in candidates:
        if os.path.isfile(p):
            weight_path = p
            break
    if weight_path is None:
        raise FileNotFoundError(f"在 {ckpt_dir} 找不到权重文件")

    logger.info(f"Loading weights from: {weight_path}")
    if weight_path.endswith(".safetensors"):
        from safetensors.torch import load_file
        state = load_file(weight_path)
    else:
        state = torch.load(weight_path, map_location="cpu")

    missing, unexpected = model.load_state_dict(state, strict=False)
    logger.info(f"load_state_dict done. missing={len(missing)} unexpected={len(unexpected)}")

# =========================
# Inference (batch)
# =========================
@torch.inference_mode()
def predict_batch(
    model,
    tokenizer,
    device: torch.device,
    input_texts: List[str],
    max_length: int,
):
    enc = tokenizer(
        input_texts,
        truncation=True,
        max_length=max_length,
        padding="longest",
        return_tensors="pt",
    )
    enc = {k: v.pin_memory().to(device, non_blocking=True) for k, v in enc.items()}

    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        out = model(**enc)

    logits = out["logits"] if isinstance(out, dict) else out.logits
    probs = torch.softmax(logits, dim=-1)
    prob_p = probs[:, 1].detach().float().cpu().tolist()
    pred_int = probs.argmax(dim=-1).detach().cpu().tolist()
    return pred_int, prob_p

# =========================
# progress (resume)
# =========================
def load_progress(progress_path: str) -> int:
    if not os.path.exists(progress_path):
        return -1
    try:
        with open(progress_path, "r", encoding="utf-8") as f:
            obj = json.load(f)
        return int(obj.get("last_done_index", -1))
    except Exception:
        return -1

def save_progress(progress_path: str, last_done_index: int):
    tmp = progress_path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump({"last_done_index": int(last_done_index)}, f, ensure_ascii=False)
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, progress_path)

# =========================
# worker main
# =========================
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base_model_name", default="openai")
    ap.add_argument("--cache_dir", default=".")
    ap.add_argument("--ckpt_dir", required=True)
    ap.add_argument("--data_jsonl", required=True)
    ap.add_argument("--out_dir", required=True)

    ap.add_argument("--batch_size", type=int, default=64)
    ap.add_argument("--max_length", type=int, default=512)

    ap.add_argument("--num_shards", type=int, required=True)
    ap.add_argument("--shard_id", type=int, required=True)

    ap.add_argument("--eval_every_steps", type=int, default=100)  # Compatible with launcher
    ap.add_argument("--flush_every_steps", type=int, default=1)  # 1=flush every step; can be larger for speed
    args = ap.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA not available: multi-GPU inference requires GPU")
    device = torch.device("cuda")

    # tokenizer
    tokenizer = AutoTokenizer.from_pretrained(
        args.base_model_name,
        trust_remote_code=True,
        cache_dir=args.cache_dir,
        local_files_only=False,
        enable_thinking=False,
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    def apply_chat_template(messages: List[Dict[str, str]]) -> str:
        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)

    # load data
    logger.info(f"[shard {args.shard_id}/{args.num_shards}] Loading data...")
    messages_all = load_messages_from_jsonl(args.data_jsonl)
    n = len(messages_all)
    logger.info(f"[shard {args.shard_id}] total samples={n}")

    # outputs per shard
    shard_tag = f"shard{args.shard_id:03d}_of_{args.num_shards:03d}"
    pred_all_path = os.path.join(args.out_dir, f"preds_all_{shard_tag}.jsonl")
    pred_err_path = os.path.join(args.out_dir, f"preds_errors_{shard_tag}.jsonl")
    table_csv_path = os.path.join(args.out_dir, f"preds_table_{shard_tag}.csv")
    progress_path = os.path.join(args.out_dir, f"progress_{shard_tag}.json")

    # resume
    last_done = load_progress(progress_path)  # Original index
    logger.info(f"[shard {args.shard_id}] resume last_done_index={last_done}")

    # open files append
    f_all = open(pred_all_path, "a", encoding="utf-8")
    f_err = open(pred_err_path, "a", encoding="utf-8")

    # CSV: Do not extract key values, just keep the fields you want + raw data
    need_header = not os.path.exists(table_csv_path) or os.path.getsize(table_csv_path) == 0
    f_csv = open(table_csv_path, "a", encoding="utf-8-sig", newline="")
    fieldnames = [
        "index",
        "gold",
        "pred",
        "pN",
        "pP",
        "prob_P",
        "user_raw",
        "input_text",
    ]
    w = csv.DictWriter(f_csv, fieldnames=fieldnames)
    if need_header:
        w.writeheader()
        f_csv.flush()
        os.fsync(f_csv.fileno())

    # model
    logger.info(f"[shard {args.shard_id}] Building model on {device}...")
    model = build_model_with_lora(args.base_model_name, args.cache_dir, device)
    model.eval()
    logger.info(f"[shard {args.shard_id}] Loading ckpt...")
    load_trained_weights(model, args.ckpt_dir)
    logger.info(f"[shard {args.shard_id}] ready.")

    # running accuracy
    gold_seen = 0
    correct = 0

    # iterate indices belonging to this shard
    def belongs(i: int) -> bool:
        return (i % args.num_shards) == args.shard_id and i > last_done

    indices = [i for i in range(n) if belongs(i)]
    total_local = len(indices)
    steps = math.ceil(total_local / args.batch_size)
    logger.info(f"[shard {args.shard_id}] local_samples={total_local}, steps={steps}, bs={args.batch_size}")

    try:
        pbar = tqdm(range(steps), desc=f"Infer {shard_tag}")
        for step in pbar:
            batch_idxs = indices[step * args.batch_size: (step + 1) * args.batch_size]

            # build batch texts
            input_texts = []
            user_raws = []
            golds = []

            for orig_i in batch_idxs:
                msgs = messages_all[orig_i]
                user_raw, gold_raw = get_user_before_last_assistant(msgs)

                # Continue using your original standardize_messages to ensure user content is a string
                std_msgs, gold_std = standardize_messages(msgs)
                gold = gold_std if gold_std else gold_raw

                input_text = apply_chat_template(std_msgs)

                input_texts.append(input_text)
                user_raws.append(user_raw if user_raw is not None else "")
                golds.append(gold)

            pred_int, prob_p_list = predict_batch(
                model=model,
                tokenizer=tokenizer,
                device=device,
                input_texts=input_texts,
                max_length=args.max_length,
            )

            # write results immediately
            for j, orig_i in enumerate(batch_idxs):
                pred = "P" if int(pred_int[j]) == 1 else "N"
                prob_p = float(prob_p_list[j])
                user_raw = user_raws[j]
                gold = golds[j]

                # CSV: Do not extract key values, 
                # just keep the fields you want + raw data
                row = {
                    "index": orig_i,
                    "gold": gold,
                    "pred": pred,
                    "pN": f"{(1.0 - prob_p):.8f}",
                    "pP": f"{prob_p:.8f}",
                    "prob_P": f"{prob_p:.8f}",
                    "user_raw": user_raw,
                    "input_text": input_texts[j],
                }
                w.writerow(row)

                # jsonl only if gold exists
                if gold in ("P", "N"):
                    item = {
                        "index": orig_i,
                        "gold": gold,
                        "pred": pred,
                        "prob_P": prob_p,
                        "user_raw": user_raw,
                        "input_text": input_texts[j],
                    }
                    f_all.write(json.dumps(item, ensure_ascii=False) + "\n")
                    if pred != gold:
                        f_err.write(json.dumps(item, ensure_ascii=False) + "\n")

                    gold_seen += 1
                    if pred == gold:
                        correct += 1

            # progress: last done = max index processed in this shard so far
            save_progress(progress_path, max(batch_idxs))

            # flush policy
            if (step + 1) % args.flush_every_steps == 0:
                f_all.flush(); os.fsync(f_all.fileno())
                f_err.flush(); os.fsync(f_err.fileno())
                f_csv.flush(); os.fsync(f_csv.fileno())

            # eval policy (Keep your original parameters and output)
            if (step + 1) % args.eval_every_steps == 0:
                if gold_seen > 0:
                    acc = correct / gold_seen
                    logger.info(f"[{shard_tag} step {step+1}/{steps}] running_acc={acc:.6f} (gold_seen={gold_seen})")
                else:
                    logger.info(f"[{shard_tag} step {step+1}/{steps}] running_acc=N/A (no gold)")

    finally:
        try:
            f_all.close()
        except Exception:
            pass
        try:
            f_err.close()
        except Exception:
            pass
        try:
            f_csv.close()
        except Exception:
            pass

    logger.info(f"[shard {args.shard_id}] DONE. outputs: {pred_all_path} {pred_err_path} {table_csv_path}")

if __name__ == "__main__":
    main()

	import os
	import json
	import math
	import time
	import csv
	import logging
	import argparse
	from typing import List, Dict, Any, Tuple, Optional

	import torch
	import torch.nn as nn
	from tqdm import tqdm

	from transformers import AutoTokenizer, AutoModel
	from peft import LoraConfig, get_peft_model

	# Flash SDP (Optional)
	torch.backends.cuda.enable_flash_sdp(True)
	torch.backends.cuda.enable_mem_efficient_sdp(True)
	torch.backends.cuda.enable_math_sdp(False)

	# =========================
	# logging
	# =========================
	logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
	logger = logging.getLogger("infer_worker")

	# =========================
	# System Prompt
	# =========================
	SYSTEM_PROMPT = """
	Act as an expert in reticular chemistry. You will receive reaction conditions as a JSON object with the fields: metal_precursor,
	organic_linker, modulator, solvent, metal_concentration_mM, M_L_ratio, temperature_C, and time_h. Based on these inputs, output exactly one uppercase label:
	'P' if the conditions are likely to yield a crystalline metal-organic framework under experimental conditions, or 'N' if not..
	"""

	# =========================
	# Data Loading
	# =========================
	def load_messages_from_jsonl(path: str) -> List[List[Dict[str, Any]]]:
	all_messages = []
	with open(path, "r", encoding="utf-8-sig") as f:
	for line in f:
	line = line.strip()
	if not line:
	continue
	data = json.loads(line)
	if "messages" in data:
	all_messages.append(data["messages"])
	return all_messages

	def standardize_json_input(json_string: str) -> str:
	try:
	data = json.loads(json_string)
	keep = {
	"metal_precursor": data.get("metal_precursor"),
	"organic_linker": data.get("organic_linker"),
	"modulator": data.get("modulator"),
	"solvent": data.get("solvent"),
	"metal_concentration_mM": data.get("metal_concentration_mM"),
	"M_L_ratio": data.get("M_L_ratio"),
	"temperature_C": data.get("temperature_C"),
	"time_h": data.get("time_h"),
	}
	return json.dumps(keep, ensure_ascii=False)
	except Exception:
	return json_string

	def char_norm(s: Optional[str]) -> Optional[str]:
	s = (s or "").strip().upper()
	return s if s in ("P", "N") else None

	def standardize_messages(messages: List[Dict[str, Any]]) -> Tuple[List[Dict[str, str]], Optional[str]]:
	user_content = None
	assistant_content = None
	for m in messages:
	if m.get("role") == "user":
	user_content = m.get("content")
	elif m.get("role") == "assistant":
	assistant_content = m.get("content")

	user_content = standardize_json_input(user_content or "")
	standardized = [
	{"role": "system", "content": SYSTEM_PROMPT},
	{"role": "user", "content": user_content},
	]
	gold = char_norm(assistant_content)
	return standardized, gold

	def get_user_before_last_assistant(messages):
	if not messages:
	return None, None
	last_asst_idx = None
	for i in range(len(messages) - 1, -1, -1):
	if messages[i].get("role") == "assistant":
	last_asst_idx = i
	break
	if last_asst_idx is None:
	return None, None
	gold_char = char_norm(messages[last_asst_idx].get("content"))
	user_raw = None
	for j in range(last_asst_idx - 1, -1, -1):
	if messages[j].get("role") == "user":
	user_raw = messages[j].get("content")
	break
	return user_raw, gold_char

	# =========================
	# Model
	# =========================
	def build_model_with_lora(base_model_name: str, cache_dir: str, device: torch.device):
	lora_config = LoraConfig(
	r=8,
	lora_alpha=16,
	target_modules=["q_proj", "v_proj", "o_proj"],
	lora_dropout=0.01,
	bias="none",
	task_type="SEQ_CLS"
	)

	backbone = AutoModel.from_pretrained(
	base_model_name,
	torch_dtype=torch.bfloat16,
	device_map=None, # 强制单卡
	cache_dir=cache_dir,
	).to(device)

	backbone = get_peft_model(backbone, lora_config)

	class QwenWithLoRAForBinaryClassification(nn.Module):
	def __init__(self, backbone):
	super().__init__()
	self.backbone = backbone
	self.config = backbone.config
	hidden_size = backbone.config.hidden_size
	self.classifier = nn.Linear(hidden_size, 2)

	def forward(self, input_ids, attention_mask=None, labels=None):
	outputs = self.backbone(
	input_ids=input_ids,
	attention_mask=attention_mask,
	return_dict=True,
	)
	mask = attention_mask.unsqueeze(-1)
	pooled = (outputs.last_hidden_state * mask).sum(1) / mask.sum(1)

	if (self.classifier.weight.device != pooled.device) or (self.classifier.weight.dtype != pooled.dtype):
	self.classifier = self.classifier.to(device=pooled.device, dtype=pooled.dtype)

	logits = self.classifier(pooled)
	loss = None
	if labels is not None:
	loss = nn.CrossEntropyLoss()(logits, labels)
	return {"loss": loss, "logits": logits}

	return QwenWithLoRAForBinaryClassification(backbone).to(device)

	def load_trained_weights(model: nn.Module, ckpt_dir: str):
	candidates = [
	os.path.join(ckpt_dir, "pytorch_model.bin"),
	os.path.join(ckpt_dir, "model.safetensors"),
	os.path.join(ckpt_dir, "adapter_model.bin"),
	]
	weight_path = None
	for p in candidates:
	if os.path.isfile(p):
	weight_path = p
	break
	if weight_path is None:
	raise FileNotFoundError(f"在 {ckpt_dir} 找不到权重文件")

	logger.info(f"Loading weights from: {weight_path}")
	if weight_path.endswith(".safetensors"):
	from safetensors.torch import load_file
	state = load_file(weight_path)
	else:
	state = torch.load(weight_path, map_location="cpu")

	missing, unexpected = model.load_state_dict(state, strict=False)
	logger.info(f"load_state_dict done. missing={len(missing)} unexpected={len(unexpected)}")

	# =========================
	# Inference (batch)
	# =========================
	@torch.inference_mode()
	def predict_batch(
	model,
	tokenizer,
	device: torch.device,
	input_texts: List[str],
	max_length: int,
	):
	enc = tokenizer(
	input_texts,
	truncation=True,
	max_length=max_length,
	padding="longest",
	return_tensors="pt",
	)
	enc = {k: v.pin_memory().to(device, non_blocking=True) for k, v in enc.items()}

	with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
	out = model(**enc)

	logits = out["logits"] if isinstance(out, dict) else out.logits
	probs = torch.softmax(logits, dim=-1)
	prob_p = probs[:, 1].detach().float().cpu().tolist()
	pred_int = probs.argmax(dim=-1).detach().cpu().tolist()
	return pred_int, prob_p

	# =========================
	# progress (resume)
	# =========================
	def load_progress(progress_path: str) -> int:
	if not os.path.exists(progress_path):
	return -1
	try:
	with open(progress_path, "r", encoding="utf-8") as f:
	obj = json.load(f)
	return int(obj.get("last_done_index", -1))
	except Exception:
	return -1

	def save_progress(progress_path: str, last_done_index: int):
	tmp = progress_path + ".tmp"
	with open(tmp, "w", encoding="utf-8") as f:
	json.dump({"last_done_index": int(last_done_index)}, f, ensure_ascii=False)
	f.flush()
	os.fsync(f.fileno())
	os.replace(tmp, progress_path)

	# =========================
	# worker main
	# =========================
	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--base_model_name", default="openai")
	ap.add_argument("--cache_dir", default=".")
	ap.add_argument("--ckpt_dir", required=True)
	ap.add_argument("--data_jsonl", required=True)
	ap.add_argument("--out_dir", required=True)

	ap.add_argument("--batch_size", type=int, default=64)
	ap.add_argument("--max_length", type=int, default=512)

	ap.add_argument("--num_shards", type=int, required=True)
	ap.add_argument("--shard_id", type=int, required=True)

	ap.add_argument("--eval_every_steps", type=int, default=100) # Compatible with launcher
	ap.add_argument("--flush_every_steps", type=int, default=1) # 1=flush every step; can be larger for speed
	args = ap.parse_args()

	os.makedirs(args.out_dir, exist_ok=True)

	if not torch.cuda.is_available():
	raise RuntimeError("CUDA not available: multi-GPU inference requires GPU")
	device = torch.device("cuda")

	# tokenizer
	tokenizer = AutoTokenizer.from_pretrained(
	args.base_model_name,
	trust_remote_code=True,
	cache_dir=args.cache_dir,
	local_files_only=False,
	enable_thinking=False,
	)
	if tokenizer.pad_token is None:
	tokenizer.pad_token = tokenizer.eos_token

	def apply_chat_template(messages: List[Dict[str, str]]) -> str:
	return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)

	# load data
	logger.info(f"[shard {args.shard_id}/{args.num_shards}] Loading data...")
	messages_all = load_messages_from_jsonl(args.data_jsonl)
	n = len(messages_all)
	logger.info(f"[shard {args.shard_id}] total samples={n}")

	# outputs per shard
	shard_tag = f"shard{args.shard_id:03d}_of_{args.num_shards:03d}"
	pred_all_path = os.path.join(args.out_dir, f"preds_all_{shard_tag}.jsonl")
	pred_err_path = os.path.join(args.out_dir, f"preds_errors_{shard_tag}.jsonl")
	table_csv_path = os.path.join(args.out_dir, f"preds_table_{shard_tag}.csv")
	progress_path = os.path.join(args.out_dir, f"progress_{shard_tag}.json")

	# resume
	last_done = load_progress(progress_path) # Original index
	logger.info(f"[shard {args.shard_id}] resume last_done_index={last_done}")

	# open files append
	f_all = open(pred_all_path, "a", encoding="utf-8")
	f_err = open(pred_err_path, "a", encoding="utf-8")

	# CSV: Do not extract key values, just keep the fields you want + raw data
	need_header = not os.path.exists(table_csv_path) or os.path.getsize(table_csv_path) == 0
	f_csv = open(table_csv_path, "a", encoding="utf-8-sig", newline="")
	fieldnames = [
	"index",
	"gold",
	"pred",
	"pN",
	"pP",
	"prob_P",
	"user_raw",
	"input_text",
	]
	w = csv.DictWriter(f_csv, fieldnames=fieldnames)
	if need_header:
	w.writeheader()
	f_csv.flush()
	os.fsync(f_csv.fileno())

	# model
	logger.info(f"[shard {args.shard_id}] Building model on {device}...")
	model = build_model_with_lora(args.base_model_name, args.cache_dir, device)
	model.eval()
	logger.info(f"[shard {args.shard_id}] Loading ckpt...")
	load_trained_weights(model, args.ckpt_dir)
	logger.info(f"[shard {args.shard_id}] ready.")

	# running accuracy
	gold_seen = 0
	correct = 0

	# iterate indices belonging to this shard
	def belongs(i: int) -> bool:
	return (i % args.num_shards) == args.shard_id and i > last_done

	indices = [i for i in range(n) if belongs(i)]
	total_local = len(indices)
	steps = math.ceil(total_local / args.batch_size)
	logger.info(f"[shard {args.shard_id}] local_samples={total_local}, steps={steps}, bs={args.batch_size}")

	try:
	pbar = tqdm(range(steps), desc=f"Infer {shard_tag}")
	for step in pbar:
	batch_idxs = indices[step * args.batch_size: (step + 1) * args.batch_size]

	# build batch texts
	input_texts = []
	user_raws = []
	golds = []

	for orig_i in batch_idxs:
	msgs = messages_all[orig_i]
	user_raw, gold_raw = get_user_before_last_assistant(msgs)

	# Continue using your original standardize_messages to ensure user content is a string
	std_msgs, gold_std = standardize_messages(msgs)
	gold = gold_std if gold_std else gold_raw

	input_text = apply_chat_template(std_msgs)

	input_texts.append(input_text)
	user_raws.append(user_raw if user_raw is not None else "")
	golds.append(gold)

	pred_int, prob_p_list = predict_batch(
	model=model,
	tokenizer=tokenizer,
	device=device,
	input_texts=input_texts,
	max_length=args.max_length,
	)

	# write results immediately
	for j, orig_i in enumerate(batch_idxs):
	pred = "P" if int(pred_int[j]) == 1 else "N"
	prob_p = float(prob_p_list[j])
	user_raw = user_raws[j]
	gold = golds[j]

	# CSV: Do not extract key values,
	# just keep the fields you want + raw data
	row = {
	"index": orig_i,
	"gold": gold,
	"pred": pred,
	"pN": f"{(1.0 - prob_p):.8f}",
	"pP": f"{prob_p:.8f}",
	"prob_P": f"{prob_p:.8f}",
	"user_raw": user_raw,
	"input_text": input_texts[j],
	}
	w.writerow(row)

	# jsonl only if gold exists
	if gold in ("P", "N"):
	item = {
	"index": orig_i,
	"gold": gold,
	"pred": pred,
	"prob_P": prob_p,
	"user_raw": user_raw,
	"input_text": input_texts[j],
	}
	f_all.write(json.dumps(item, ensure_ascii=False) + "\n")
	if pred != gold:
	f_err.write(json.dumps(item, ensure_ascii=False) + "\n")

	gold_seen += 1
	if pred == gold:
	correct += 1

	# progress: last done = max index processed in this shard so far
	save_progress(progress_path, max(batch_idxs))

	# flush policy
	if (step + 1) % args.flush_every_steps == 0:
	f_all.flush(); os.fsync(f_all.fileno())
	f_err.flush(); os.fsync(f_err.fileno())
	f_csv.flush(); os.fsync(f_csv.fileno())

	# eval policy (Keep your original parameters and output)
	if (step + 1) % args.eval_every_steps == 0:
	if gold_seen > 0:
	acc = correct / gold_seen
	logger.info(f"[{shard_tag} step {step+1}/{steps}] running_acc={acc:.6f} (gold_seen={gold_seen})")
	else:
	logger.info(f"[{shard_tag} step {step+1}/{steps}] running_acc=N/A (no gold)")

	finally:
	try:
	f_all.close()
	except Exception:
	pass
	try:
	f_err.close()
	except Exception:
	pass
	try:
	f_csv.close()
	except Exception:
	pass

	logger.info(f"[shard {args.shard_id}] DONE. outputs: {pred_all_path} {pred_err_path} {table_csv_path}")

	if __name__ == "__main__":
	main()