|
|
|
|
|
import os |
|
import re |
|
from collections import defaultdict |
|
from pathlib import Path |
|
from typing import Dict, List, Tuple |
|
|
|
from update_readme import generate_url, get_all_files |
|
|
|
|
|
class Wheel: |
|
def __init__(self, full_name: str, url: str): |
|
""" |
|
Args: |
|
full_name: |
|
Example: k2-1.23.4.dev20230223+cpu.torch1.10.0-cp36-cp36m-linux_x86_64.whl |
|
""" |
|
self.full_name = full_name |
|
|
|
pattern = ( |
|
r"k2-(\d)\.(\d)+((\.)(\d))?\.dev(\d{8})\+cpu\.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.torch_version = m.group(7) |
|
self.py_version = int(m.group(8)) |
|
self.url = url |
|
|
|
def __str__(self): |
|
return self.url |
|
|
|
def __repr__(self): |
|
return self.url |
|
|
|
|
|
def generate_index(filename: str, torch_versions) -> str: |
|
b = [] |
|
for i in torch_versions: |
|
b.append(f" ./{i}.rst") |
|
b = "\n".join(b) |
|
|
|
s = f"""\ |
|
Pre-compiled CPU wheels (Linux) |
|
=============================== |
|
|
|
This page describes pre-compiled ``CPU`` wheels for `k2`_ on Linux. |
|
|
|
.. toctree:: |
|
:maxdepth: 2 |
|
|
|
{b} |
|
""" |
|
with open(filename, "w") as f: |
|
f.write(s) |
|
|
|
|
|
def sort_by_wheel(x: Wheel): |
|
return x.k2_major, x.k2_minor, x.k2_patch, x.k2_date, x.py_version |
|
|
|
|
|
def sort_by_torch(x): |
|
major, minor, patch = x.split(".") |
|
return int(major), int(minor), int(patch) |
|
|
|
|
|
def get_all_torch_versions(wheels: List[Wheel]) -> List[str]: |
|
ans = set() |
|
for w in wheels: |
|
ans.add(w.torch_version) |
|
|
|
|
|
ans = list(ans) |
|
ans.sort(reverse=True, key=sort_by_torch) |
|
return ans |
|
|
|
|
|
def get_doc_dir(): |
|
k2_dir = os.getenv("K2_DIR") |
|
if k2_dir is None: |
|
raise ValueError("Please set the environment variable k2_dir") |
|
|
|
cpu_dir = Path(k2_dir) / "docs/source/installation/pre-compiled-cpu-wheels-linux" |
|
|
|
if not Path(cpu_dir).is_dir(): |
|
raise ValueError(f"{cpu_dir} does not exist") |
|
|
|
print(f"k2 doc cpu_dir: {cpu_dir}") |
|
return cpu_dir |
|
|
|
|
|
def remove_all_files(d: str): |
|
files = get_all_files(d, "*.rst") |
|
for f in files: |
|
print(f"removing {f}") |
|
os.remove(f) |
|
|
|
|
|
def get_all_cpu_wheels(): |
|
cpu = get_all_files("cpu", suffix="*.whl") |
|
cpu_wheels = generate_url(cpu) |
|
return cpu_wheels |
|
|
|
|
|
def generate_file(d: str, torch_version: str, wheels: List[Wheel]) -> str: |
|
s = f"torch {torch_version}\n" |
|
s += "=" * len(f"torch {torch_version}") |
|
s += "\n" * 3 |
|
wheels = filter(lambda w: w.torch_version == torch_version, wheels) |
|
wheels = list(wheels) |
|
wheels.sort(reverse=True, key=sort_by_wheel) |
|
for w in wheels: |
|
s += f"- `{w.full_name} <{w.url}>`_\n" |
|
|
|
with open(f"{d}/{torch_version}.rst", "w") as f: |
|
f.write(s) |
|
|
|
|
|
def main(): |
|
d = get_doc_dir() |
|
remove_all_files(d) |
|
|
|
urls = get_all_cpu_wheels() |
|
|
|
wheels = [] |
|
for url in urls: |
|
full_name = url.rsplit("/", maxsplit=1)[1] |
|
wheels.append(Wheel(full_name, url)) |
|
torch_versions = get_all_torch_versions(wheels) |
|
|
|
content = [] |
|
for t in torch_versions: |
|
s = generate_file(d, t, wheels) |
|
|
|
generate_index(f"{d}/index.rst", torch_versions) |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
