|
|
|
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 |
|
k2-1.24.4.dev20240301+cuda12.1.torch2.3.0.dev20240229-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl |
|
""" |
|
self.full_name = str(full_name) |
|
pattern = r"k2-(\d)\.(\d)+((\.)(\d))?\.dev(\d{8})\+cuda(\d+)\.(\d+)\.torch(\d\.\d+\.\d(\.dev\d{8})?)-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(11)) |
|
self.url = url |
|
|
|
|
|
def sort_by_wheel(x: Wheel): |
|
torch_major, torch_minor, torch_patch = x.torch_version.split(".")[:3] |
|
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') |
|
with open(d / "source/cuda-cn.html", "w") as f: |
|
for w in wheels: |
|
url = w.url.replace("https://huggingface.co", "https://hf-mirror.com") |
|
f.write(f'<a href="{url}">{w.full_name}</a><br/>\n') |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
