fix: handle long sequences and optimize IFD computation

This commit is contained in:
ViperEkura 2026-07-04 08:35:45 +08:00
parent 204873fa2f
commit 1adca39cd8
2 changed files with 10 additions and 4 deletions

6
.gitignore vendored
View File

@ -5,8 +5,10 @@
!*/ !*/
# Allow specific file types and root files # Allow specific file types and root files
!*.py !astrai/**/*.py
!*.sh !scripts/**/*.py
!scripts/**/*.sh
!tests/**/*.py
# Allow GitHub files # Allow GitHub files
!/.github/** !/.github/**

View File

@ -89,8 +89,12 @@ def _compute_ifd_raw(model, tokenizer, instruction, response, device, max_len) -
qa_ids = instr_ids + resp_ids qa_ids = instr_ids + resp_ids
with torch.inference_mode(): with torch.inference_mode():
logits_qa = model(torch.tensor([qa_ids], device=device, dtype=torch.long))["logits"][0] logits_qa = model(torch.tensor([qa_ids], device=device, dtype=torch.long))[
logits_resp = model(torch.tensor([resp_ids], device=device, dtype=torch.long))["logits"][0] "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 = logits_resp.new_tensor(resp_ids, dtype=torch.long) resp_targets = logits_resp.new_tensor(resp_ids, dtype=torch.long)