File size: 2,061 Bytes
6ac7d7a 724b38e 6ac7d7a |
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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
#!/usr/bin/env python3
import os
import re
from pathlib import Path
from update_readme import generate_url, get_all_files
def get_doc_dir():
k2_dir = os.getenv("K2_DIR")
if k2_dir is None:
raise ValueError("Please set the environment variable k2_dir")
return Path(k2_dir) / "docs"
class Wheel:
def __init__(self, full_name: str, url: str):
"""
Args:
full_name:
Example: k2-1.23.4.dev20230224+cuda10.1.torch1.6.0-cp36-cp36m-linux_x86_64.whl
"""
self.full_name = full_name
pattern = r"k2-(\d)\.(\d)+((\.)(\d))?\.dev(\d{8})\+cuda(\d+)\.(\d+)\.torch(\d\.\d+\.\d)-cp(\d+)"
m = re.search(pattern, full_name)
self.k2_major = int(m.group(1))
self.k2_minor = int(m.group(2))
self.k2_patch = int(m.group(5))
self.k2_date = int(m.group(6))
self.cuda_major_version = int(m.group(7))
self.cuda_minor_version = int(m.group(8))
self.torch_version = m.group(9)
self.py_version = int(m.group(10))
self.url = url
def sort_by_wheel(x: Wheel):
torch_major, torch_minor, torch_patch = x.torch_version.split(".")
torch_major = int(torch_major)
torch_minor = int(torch_minor)
torch_patch = int(torch_patch)
return (
x.k2_major,
x.k2_minor,
x.k2_patch,
x.k2_date,
torch_major,
torch_minor,
torch_patch,
x.cuda_major_version,
x.cuda_minor_version,
x.py_version,
)
def main():
wheels = get_all_files("cuda", suffix="*.whl")
wheels += get_all_files("ubuntu-cuda", suffix="*.whl")
urls = generate_url(wheels)
wheels = []
for url in urls:
full_name = url.rsplit("/", maxsplit=1)[1]
wheels.append(Wheel(full_name, url))
wheels.sort(reverse=True, key=sort_by_wheel)
d = get_doc_dir()
with open(d / "source/cuda.html", "w") as f:
for w in wheels:
f.write(f'<a href="{w.url}">{w.full_name}</a><br/>\n')
if __name__ == "__main__":
main()
|