refactor: rewrite IFD evaluation with clean three-layer architecture
This commit is contained in:
parent
4e508afa2d
commit
4d3c9341c1
|
|
@ -1,31 +1,14 @@
|
|||
"""IFD (Instruction Following Difficulty) data quality scoring.
|
||||
|
||||
Computes IFD scores for instruction-response pairs to guide data selection.
|
||||
IFD = conditional_NLL / unconditional_NLL, where:
|
||||
IFD = conditional_NLL / unconditional_NLL
|
||||
|
||||
- conditional_NLL: average CE loss on response tokens given instruction context
|
||||
- unconditional_NLL: average CE loss on response tokens alone
|
||||
|
||||
Higher IFD (close to 1) = instruction provides less help = harder sample.
|
||||
Lower IFD (close to 0) = instruction provides strong guidance = easy sample.
|
||||
IFD > 1 = instruction misleads the model = likely low-quality data.
|
||||
|
||||
Usage::
|
||||
|
||||
python scripts/eval/ifd.py --param_path ./params \
|
||||
--input data.jsonl --output data_with_ifd.jsonl \
|
||||
--instr_key instruction --resp_key response
|
||||
|
||||
Disable chat template::
|
||||
|
||||
python scripts/eval/ifd.py --param_path ./params \
|
||||
--input data.jsonl --output data_with_ifd.jsonl \
|
||||
--instr_key instruction --resp_key response \
|
||||
--no_chat_template
|
||||
- Messages format: plain text concatenation (no chat template)
|
||||
- Plain format: raw instr_key + resp_key fields
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import statistics
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
|
@ -35,172 +18,88 @@ from astrai.model import AutoModel
|
|||
from astrai.tokenize import AutoTokenizer
|
||||
|
||||
|
||||
def compute_ifd(
|
||||
model,
|
||||
tokenizer,
|
||||
instruction: str,
|
||||
response: str,
|
||||
device: str,
|
||||
max_len: int = 2048,
|
||||
use_chat_template: bool = False,
|
||||
) -> dict:
|
||||
if use_chat_template:
|
||||
return _compute_ifd_with_template(
|
||||
model, tokenizer, instruction, response, device, max_len
|
||||
)
|
||||
return _compute_ifd_raw(model, tokenizer, instruction, response, device, max_len)
|
||||
|
||||
|
||||
def _compute_ifd_raw(model, tokenizer, instruction, response, device, max_len) -> dict:
|
||||
instr_ids = tokenizer.encode(instruction, add_special_tokens=False)
|
||||
resp_ids = tokenizer.encode(response, add_special_tokens=False)
|
||||
|
||||
if len(resp_ids) > max_len:
|
||||
resp_ids = resp_ids[:max_len]
|
||||
|
||||
def _score(context_ids, resp_ids, model, device):
|
||||
"""Core IFD computation: context → L_cond, response alone → L_uncond."""
|
||||
if not resp_ids:
|
||||
return {
|
||||
"L_cond": None,
|
||||
"L_uncond": None,
|
||||
"ifd": None,
|
||||
"error": "empty response",
|
||||
}
|
||||
|
||||
qa_len = len(instr_ids) + len(resp_ids)
|
||||
if qa_len > max_len:
|
||||
overflow = qa_len - max_len
|
||||
if overflow >= len(instr_ids):
|
||||
resp_ids = resp_ids[:max_len]
|
||||
instr_ids = []
|
||||
else:
|
||||
instr_ids = instr_ids[overflow:]
|
||||
|
||||
if not instr_ids:
|
||||
return {
|
||||
"L_cond": None,
|
||||
"L_uncond": None,
|
||||
"ifd": None,
|
||||
"error": "response too long for context",
|
||||
}
|
||||
|
||||
instr_len = len(instr_ids)
|
||||
resp_len = len(resp_ids)
|
||||
|
||||
qa_ids = instr_ids + resp_ids
|
||||
|
||||
with torch.inference_mode():
|
||||
logits_qa = model(torch.tensor([qa_ids], device=device, dtype=torch.long))[
|
||||
"logits"
|
||||
][0]
|
||||
logits_resp = model(torch.tensor([resp_ids], device=device, dtype=torch.long))[
|
||||
"logits"
|
||||
][0]
|
||||
|
||||
resp_logits = logits_qa[instr_len - 1 : -1]
|
||||
resp_targets = logits_resp.new_tensor(resp_ids, dtype=torch.long)
|
||||
return None
|
||||
full_ids = context_ids + resp_ids
|
||||
inp_full = torch.tensor([full_ids], device=device, dtype=torch.long)
|
||||
inp_resp = torch.tensor([resp_ids], device=device, dtype=torch.long)
|
||||
logits_full = model(inp_full)["logits"][0]
|
||||
logits_resp = model(inp_resp)["logits"][0]
|
||||
ctx_len = len(context_ids)
|
||||
resp_logits = logits_full[ctx_len - 1 : -1]
|
||||
resp_targets = torch.tensor(resp_ids, device=device, dtype=torch.long)
|
||||
L_cond = F.cross_entropy(resp_logits, resp_targets, reduction="mean").item()
|
||||
|
||||
unp_logits = logits_resp[:-1]
|
||||
unp_targets = logits_resp.new_tensor(resp_ids[1:], dtype=torch.long)
|
||||
unp_targets = torch.tensor(resp_ids[1:], device=device, dtype=torch.long)
|
||||
L_uncond = F.cross_entropy(unp_logits, unp_targets, reduction="mean").item()
|
||||
|
||||
ifd = L_cond / L_uncond if L_uncond > 0 else None
|
||||
|
||||
return {
|
||||
"L_cond": round(L_cond, 6),
|
||||
"L_uncond": round(L_uncond, 6),
|
||||
"ifd": round(ifd, 6) if ifd is not None else None,
|
||||
"instr_len": instr_len,
|
||||
"resp_len": resp_len,
|
||||
"error": None,
|
||||
}
|
||||
|
||||
|
||||
def _compute_ifd_with_template(
|
||||
model, tokenizer, instruction, response, device, max_len
|
||||
) -> dict:
|
||||
instr_prefix = tokenizer.apply_chat_template(
|
||||
[{"role": "user", "content": instruction}],
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
)
|
||||
full_text = tokenizer.apply_chat_template(
|
||||
[
|
||||
{"role": "user", "content": instruction},
|
||||
{"role": "assistant", "content": response},
|
||||
],
|
||||
tokenize=False,
|
||||
add_generation_prompt=False,
|
||||
)
|
||||
|
||||
full_ids = tokenizer.encode(full_text)
|
||||
prefix_ids = tokenizer.encode(instr_prefix)
|
||||
resp_ids = tokenizer.encode(response)
|
||||
|
||||
if not resp_ids:
|
||||
return {
|
||||
"L_cond": None,
|
||||
"L_uncond": None,
|
||||
"ifd": None,
|
||||
"error": "empty response",
|
||||
}
|
||||
|
||||
if len(full_ids) > max_len:
|
||||
overflow = len(full_ids) - max_len
|
||||
full_ids = full_ids[overflow:]
|
||||
prefix_len = len(prefix_ids) - overflow
|
||||
prefix_len = max(0, prefix_len)
|
||||
else:
|
||||
prefix_len = len(prefix_ids)
|
||||
|
||||
cond_tensor = torch.tensor([full_ids], device=device, dtype=torch.long)
|
||||
|
||||
with torch.inference_mode():
|
||||
logits_qa = model(cond_tensor)["logits"][0]
|
||||
|
||||
resp_start = prefix_len - 1
|
||||
resp_end = len(full_ids) - 1
|
||||
if resp_end <= resp_start:
|
||||
return {
|
||||
"L_cond": None,
|
||||
"L_uncond": None,
|
||||
"ifd": None,
|
||||
"error": "response truncated entirely",
|
||||
}
|
||||
|
||||
resp_logits = logits_qa[resp_start:resp_end]
|
||||
resp_targets = torch.tensor(full_ids[prefix_len:], device=device, dtype=torch.long)
|
||||
L_cond = F.cross_entropy(resp_logits, resp_targets, reduction="mean").item()
|
||||
|
||||
resp_tensor = torch.tensor([resp_ids], device=device, dtype=torch.long)
|
||||
|
||||
with torch.inference_mode():
|
||||
logits_resp = model(resp_tensor)["logits"][0]
|
||||
|
||||
unp_logits = logits_resp[:-1]
|
||||
unp_targets = resp_tensor[0, 1:]
|
||||
L_uncond = F.cross_entropy(unp_logits, unp_targets, reduction="mean").item()
|
||||
|
||||
ifd = L_cond / L_uncond if L_uncond > 0 else None
|
||||
|
||||
return {
|
||||
"L_cond": round(L_cond, 6),
|
||||
"L_uncond": round(L_uncond, 6),
|
||||
"ifd": round(ifd, 6) if ifd is not None else None,
|
||||
"instr_len": prefix_len,
|
||||
"resp_len": len(resp_ids),
|
||||
"error": None,
|
||||
}
|
||||
|
||||
|
||||
def _trim(context_ids, resp_ids, max_len):
|
||||
"""Truncate to fit max_len, keeping response intact if possible."""
|
||||
if len(resp_ids) > max_len // 2:
|
||||
resp_ids = resp_ids[: max_len // 2]
|
||||
full_ids = context_ids + resp_ids
|
||||
if len(full_ids) <= max_len:
|
||||
return context_ids, resp_ids
|
||||
overflow = len(full_ids) - max_len
|
||||
if overflow >= len(context_ids):
|
||||
return [], resp_ids[:max_len]
|
||||
return context_ids[overflow:], resp_ids
|
||||
|
||||
|
||||
def score_plain(model, tokenizer, instruction, response, device, max_len=2048):
|
||||
"""Compute IFD for a single instruction-response pair (plain format)."""
|
||||
ctx_ids = tokenizer.encode(instruction, add_special_tokens=False)
|
||||
resp_ids = tokenizer.encode(response, add_special_tokens=False)
|
||||
ctx_ids, resp_ids = _trim(ctx_ids, resp_ids, max_len)
|
||||
if not ctx_ids or not resp_ids:
|
||||
return {"L_cond": None, "L_uncond": None, "ifd": None, "error": "empty"}
|
||||
return _score(ctx_ids, resp_ids, model, device)
|
||||
|
||||
|
||||
def score_messages(model, tokenizer, messages, device, max_len=2048):
|
||||
"""Compute IFD for each assistant turn in a messages array."""
|
||||
turns = []
|
||||
for i, msg in enumerate(messages):
|
||||
if msg.get("role") != "assistant":
|
||||
continue
|
||||
ctx_text = "\n\n".join(m["content"] for m in messages[:i])
|
||||
ctx_ids = tokenizer.encode(ctx_text)
|
||||
resp_ids = tokenizer.encode(msg["content"], add_special_tokens=False)
|
||||
ctx_ids, resp_ids = _trim(ctx_ids, resp_ids, max_len)
|
||||
if ctx_ids and resp_ids:
|
||||
turns.append(_score(ctx_ids, resp_ids, model, device))
|
||||
if not turns:
|
||||
return None
|
||||
valid = [t for t in turns if t is not None and t["ifd"] is not None]
|
||||
if not valid:
|
||||
return {"ifd": None, "ifd_turns": turns}
|
||||
avg = sum(t["ifd"] for t in valid) / len(valid)
|
||||
return {
|
||||
"ifd": avg,
|
||||
"ifd_detail": valid[0] if len(valid) == 1 else None,
|
||||
"ifd_turns": turns,
|
||||
}
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def process_file(
|
||||
param_path: str,
|
||||
input_file: str,
|
||||
output_file: str,
|
||||
instr_key: str,
|
||||
resp_key: str,
|
||||
max_len: int = 2048,
|
||||
use_chat_template: bool = False,
|
||||
param_path,
|
||||
input_file,
|
||||
output_file,
|
||||
instr_key,
|
||||
resp_key,
|
||||
max_len=2048,
|
||||
data_format="plain",
|
||||
):
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
dtype = torch.bfloat16 if device == "cuda" else torch.float32
|
||||
|
|
@ -210,42 +109,35 @@ def process_file(
|
|||
model.to(device=device, dtype=dtype)
|
||||
model.eval()
|
||||
|
||||
if use_chat_template and tokenizer._chat_template is None:
|
||||
raise RuntimeError(
|
||||
"--use_chat_template specified but tokenizer has no chat template. "
|
||||
"Add a chat_template to tokenizer_config.json or omit the flag."
|
||||
)
|
||||
|
||||
with open(input_file, "r", encoding="utf-8") as f:
|
||||
with open(input_file, encoding="utf-8") as f:
|
||||
data = [json.loads(line) for line in f if line.strip()]
|
||||
|
||||
results = []
|
||||
ifd_values = []
|
||||
all_ifds = []
|
||||
|
||||
with torch.inference_mode():
|
||||
for item in tqdm.tqdm(data, desc="Computing IFD", unit="sample"):
|
||||
instruction = item[instr_key]
|
||||
response = item[resp_key]
|
||||
scores = compute_ifd(
|
||||
model,
|
||||
tokenizer,
|
||||
instruction,
|
||||
response,
|
||||
device,
|
||||
max_len,
|
||||
use_chat_template=use_chat_template,
|
||||
for item in tqdm.tqdm(data, desc="Computing IFD", unit="sample"):
|
||||
if data_format == "messages":
|
||||
scores = score_messages(
|
||||
model, tokenizer, item.get("messages", []), device, max_len
|
||||
)
|
||||
ifd_values.append(scores["ifd"])
|
||||
if scores is None:
|
||||
results.append({**item, "ifd": None, "ifd_turns": []})
|
||||
else:
|
||||
all_ifds.append(scores["ifd"])
|
||||
results.append({**item, **scores})
|
||||
else:
|
||||
scores = score_plain(
|
||||
model, tokenizer, item[instr_key], item[resp_key], device, max_len
|
||||
)
|
||||
all_ifds.append(scores["ifd"])
|
||||
results.append({**item, "ifd": scores["ifd"], "ifd_detail": scores})
|
||||
|
||||
with open(output_file, "w", encoding="utf-8") as f:
|
||||
for item in results:
|
||||
f.write(json.dumps(item, ensure_ascii=False) + "\n")
|
||||
|
||||
valid_ifd = [v for v in ifd_values if v is not None]
|
||||
valid_ifd = [v for v in all_ifds if v is not None]
|
||||
if valid_ifd:
|
||||
import statistics
|
||||
|
||||
print(f"\n{'=' * 50}")
|
||||
print(f" Samples: {len(data)}")
|
||||
print(f" Valid IFD: {len(valid_ifd)}")
|
||||
|
|
@ -266,29 +158,19 @@ def main():
|
|||
parser.add_argument("--param_path", type=str, required=True, help="Model directory")
|
||||
parser.add_argument("--input", type=str, required=True, help="Input JSONL file")
|
||||
parser.add_argument("--output", type=str, required=True, help="Output JSONL file")
|
||||
parser.add_argument("--max_len", type=int, default=2048, help="Max token length")
|
||||
parser.add_argument(
|
||||
"--instr_key",
|
||||
"--format",
|
||||
type=str,
|
||||
default="instruction",
|
||||
help="Key for instruction field",
|
||||
default="plain",
|
||||
choices=["plain", "messages"],
|
||||
help="Input format: 'plain' for instr_key+resp_key, 'messages' for messages array",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--resp_key",
|
||||
type=str,
|
||||
default="response",
|
||||
help="Key for response field",
|
||||
"--instr_key", type=str, default="instruction", help="Key for instruction field"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_len",
|
||||
type=int,
|
||||
default=2048,
|
||||
help="Max token length (instruction truncated to fit)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no_chat_template",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Disable chat template, use raw text concatenation",
|
||||
"--resp_key", type=str, default="response", help="Key for response field"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
|
|
@ -299,7 +181,7 @@ def main():
|
|||
args.instr_key,
|
||||
args.resp_key,
|
||||
args.max_len,
|
||||
use_chat_template=not args.no_chat_template,
|
||||
data_format=args.format,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue