ViperEkura
/
AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: handle long sequences and optimize IFD computation

This commit is contained in:
ViperEkura 2026-07-04 08:21:53 +08:00
parent a5c1de6b1b
commit 204873fa2f
1 changed files with 23 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
scripts/eval/evaluate_ifd.py
View File

@ -52,8 +52,11 @@ def compute_ifd(
def _compute_ifd_raw(model, tokenizer, instruction, response, device, max_len) -> dict: def _compute_ifd_raw(model, tokenizer, instruction, response, device, max_len) -> dict:
instr_ids = tokenizer.encode(instruction) instr_ids = tokenizer.encode(instruction, add_special_tokens=False)
resp_ids = tokenizer.encode(response) resp_ids = tokenizer.encode(response, add_special_tokens=False)
if len(resp_ids) > max_len:
resp_ids = resp_ids[:max_len]
if not resp_ids: if not resp_ids:
return { return {
@ -66,28 +69,35 @@ def _compute_ifd_raw(model, tokenizer, instruction, response, device, max_len) -
qa_len = len(instr_ids) + len(resp_ids) qa_len = len(instr_ids) + len(resp_ids)
if qa_len > max_len: if qa_len > max_len:
overflow = 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:] 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) instr_len = len(instr_ids)
resp_len = len(resp_ids) resp_len = len(resp_ids)
qa_ids = instr_ids + resp_ids qa_ids = instr_ids + resp_ids
qa_tensor = torch.tensor([qa_ids], device=device, dtype=torch.long)
with torch.inference_mode(): with torch.inference_mode():
logits_qa = model(qa_tensor)["logits"][0] 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_logits = logits_qa[instr_len - 1 : -1]
resp_targets = torch.tensor(resp_ids, device=device, dtype=torch.long) resp_targets = logits_resp.new_tensor(resp_ids, dtype=torch.long)
L_cond = F.cross_entropy(resp_logits, resp_targets, reduction="mean").item() 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_logits = logits_resp[:-1]
unp_targets = resp_tensor[0, 1:] unp_targets = logits_resp.new_tensor(resp_ids[1:], dtype=torch.long)
L_uncond = F.cross_entropy(unp_logits, unp_targets, reduction="mean").item() L_uncond = F.cross_entropy(unp_logits, unp_targets, reduction="mean").item()
ifd = L_cond / L_uncond if L_uncond > 0 else None ifd = L_cond / L_uncond if L_uncond > 0 else None
@ -185,7 +195,7 @@ def process_file(
output_file: str, output_file: str,
instr_key: str, instr_key: str,
resp_key: str, resp_key: str,
max_len: int, max_len: int = 2048,
use_chat_template: bool = False, use_chat_template: bool = False,
): ):
device = "cuda" if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu"