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
!*.py
!*.sh
!astrai/**/*.py
!scripts/**/*.py
!scripts/**/*.sh
!tests/**/*.py
# Allow GitHub files
!/.github/**

View File

@ -89,8 +89,12 @@ def _compute_ifd_raw(model, tokenizer, instruction, response, device, max_len) -
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]
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)