1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
|
- import tempfile
- from pathlib import Path
- import numpy as np
- import pytest
- import torch
- from madewithml import utils
- def test_set_seed():
- utils.set_seeds()
- a = np.random.randn(2, 3)
- b = np.random.randn(2, 3)
- utils.set_seeds()
- x = np.random.randn(2, 3)
- y = np.random.randn(2, 3)
- assert np.array_equal(a, x)
- assert np.array_equal(b, y)
- def test_save_and_load_dict():
- with tempfile.TemporaryDirectory() as dp:
- d = {"hello": "world"}
- fp = Path(dp, "d.json")
- utils.save_dict(d=d, path=fp)
- d = utils.load_dict(path=fp)
- assert d["hello"] == "world"
- def test_pad_array():
- arr = np.array([[1, 2], [1, 2, 3]], dtype="object")
- padded_arr = np.array([[1, 2, 0], [1, 2, 3]])
- assert np.array_equal(utils.pad_array(arr), padded_arr)
- def test_collate_fn():
- batch = {
- "ids": np.array([[1, 2], [1, 2, 3]], dtype="object"),
- "masks": np.array([[1, 1], [1, 1, 1]], dtype="object"),
- "targets": np.array([3, 1]),
- }
- processed_batch = utils.collate_fn(batch)
- expected_batch = {
- "ids": torch.tensor([[1, 2, 0], [1, 2, 3]], dtype=torch.int32),
- "masks": torch.tensor([[1, 1, 0], [1, 1, 1]], dtype=torch.int32),
- "targets": torch.tensor([3, 1], dtype=torch.int64),
- }
- for k in batch:
- assert torch.allclose(processed_batch[k], expected_batch[k])
- @pytest.mark.parametrize(
- "d, keys, list",
- [
- ({"a": [1, 2], "b": [1, 2]}, ["a", "b"], [{"a": 1, "b": 1}, {"a": 2, "b": 2}]),
- ({"a": [1, 2], "b": [1, 2]}, ["a"], [{"a": 1}, {"a": 2}]),
- ],
- )
- def test_dict_to_list(d, keys, list):
- assert utils.dict_to_list(d, keys=keys) == list
|