| 9 | |
| 10 | |
| 11 | def make_cache(location: Path | None): |
| 12 | def cache(key: Any, fallback: Callable[[], T]) -> T: |
| 13 | # If there's no cache location, the cache is disabled. |
| 14 | if location is None: |
| 15 | return fallback() |
| 16 | |
| 17 | key_str = hashlib.sha256(json.dumps(key).encode("utf-8")).digest().hex() |
| 18 | |
| 19 | path = location / f"{key_str}.torch" |
| 20 | try: |
| 21 | # Attempt to load the cached item. |
| 22 | key_loaded, value = torch.load(path) |
| 23 | |
| 24 | # If there was a hash collision and the keys don't actually match, throw an |
| 25 | # error so that the fallback can be used. |
| 26 | if key != key_loaded: |
| 27 | raise ValueError("Keys did not match!") |
| 28 | |
| 29 | return value |
| 30 | except (FileNotFoundError, ValueError): |
| 31 | # Use the fallback to compute the value. |
| 32 | value = fallback() |
| 33 | |
| 34 | # Cache the value. |
| 35 | path.parent.mkdir(exist_ok=True, parents=True) |
| 36 | torch.save((key, value), path) |
| 37 | |
| 38 | return value |
| 39 | |
| 40 | return cache |