Given any path belonging to a multi-file model (e.g. foo.bin.1), return the nth path in the model.
(path: Path, n: int)
| 1506 | |
| 1507 | |
| 1508 | def nth_multifile_path(path: Path, n: int) -> Path | None: |
| 1509 | '''Given any path belonging to a multi-file model (e.g. foo.bin.1), return |
| 1510 | the nth path in the model. |
| 1511 | ''' |
| 1512 | # Support the following patterns: |
| 1513 | patterns = [ |
| 1514 | # - x.00.pth, x.01.pth, etc. |
| 1515 | (r'\.[0-9]{2}\.pth$', f'.{n:02}.pth'), |
| 1516 | # - x-00001-of-00002.bin, x-00002-of-00002.bin, etc. |
| 1517 | (r'-[0-9]{5}-of-(.*)$', fr'-{n:05}-of-\1'), |
| 1518 | # x.bin, x.bin.1, etc. |
| 1519 | (r'(\.[0-9]+)?$', r'\1' if n == 0 else fr'\1.{n}') |
| 1520 | ] |
| 1521 | for regex, replacement in patterns: |
| 1522 | if re.search(regex, path.name): |
| 1523 | new_path = path.with_name(re.sub(regex, replacement, path.name)) |
| 1524 | if new_path.exists(): |
| 1525 | return new_path |
| 1526 | return None |
| 1527 | |
| 1528 | |
| 1529 | def find_multifile_paths(path: Path) -> list[Path]: |
no outgoing calls
no test coverage detected