Skip to content

Commit 4df56a4

Browse files
committed
OLMo 7B integration test fix
1 parent 450f038 commit 4df56a4

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

tests/models/olmo/test_modeling_olmo.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -371,10 +371,10 @@ def test_model_7b_logits(self):
371371
model = OLMoForCausalLM.from_pretrained("allenai/OLMo-7B-hf")
372372
out = model(torch.tensor(input_ids)).logits
373373
# Expected mean on dim = -1
374-
EXPECTED_MEAN = torch.tensor([[0.0271, 0.0249, -0.0578, -0.0870, 0.0167, 0.0710, 0.1002, 0.0677]])
374+
EXPECTED_MEAN = torch.tensor([[0.0266, -0.0012, -0.0589, -0.0868, 0.0284, 0.0609, 0.0836, 0.0552]])
375375
torch.testing.assert_close(out.mean(-1), EXPECTED_MEAN, atol=1e-2, rtol=1e-2)
376376
# slicing logits[0, 0, 0:30]
377-
EXPECTED_SLICE = torch.tensor([-1.7433, -1.6685, 7.4941, 6.1506, 0.1364, -0.1127, 1.3224, 4.5458, 4.2068, 5.8296, 7.4723, 2.7925, 3.1245, 10.8872, 10.0758, 10.6717, 7.0945, 1.2398, 3.6766, 4.2365, 2.5655, 2.2222, 1.7418, 0.5223, 0.7753, 1.0938, 0.6723, 6.2522, 6.2264, 1.8105]) # fmt: skip
377+
EXPECTED_SLICE = torch.tensor([-1.5364, -1.5505, 6.6963, 4.6085, -0.5580, 0.3295, 2.0886, 4.2162, 3.3459, 4.6850, 6.9190, 2.5602, 2.6904, 9.5714, 9.6181, 9.7108, 6.0352, 1.0118, 3.6572, 3.8474, 1.9998, 1.7586, 1.6130, 0.8734, 0.5737, 1.2120, 0.2762, 5.5266, 5.8129, 1.8287]) # fmt: skip
378378
torch.testing.assert_close(out[0, 0, :30], EXPECTED_SLICE, atol=1e-2, rtol=1e-2)
379379

380380
@unittest.skip("Logits are not yet correct, will update!")

0 commit comments

Comments
 (0)