small fixes

update-cuda-html.py

@@ -20,9 +20,10 @@ class Wheel:
         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 = full_name
-        pattern = r"k2-(\d)\.(\d)+((\.)(\d))?\.dev(\d{8})\+cuda(\d+)\.(\d+)\.torch(\d\.\d+\.\d)-cp(\d+)"
+        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))
@@ -32,12 +33,12 @@ class Wheel:
         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.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(".")
+    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)

update-k2-doc-cuda-linux.py

@@ -15,9 +15,10 @@ class Wheel:
         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 = full_name
-        pattern = r"k2-(\d)\.(\d)+((\.)(\d))?\.dev(\d{8})\+cuda(\d+)\.(\d+)\.torch(\d\.\d+\.\d)-cp(\d+)"
+        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))
@@ -27,7 +28,7 @@ class Wheel:
         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.py_version = int(m.group(11))
         self.url = url
 
     def __str__(self):
@@ -78,7 +79,7 @@ def sort_by_torch(x):
 def get_all_torch_versions(wheels: List[Wheel]) -> List[str]:
     ans = set()
     for w in wheels:
-        ans.add(w.torch_version)
+        ans.add(".".join(w.torch_version.split(".")[:3]))
 
     # sort torch version from high to low
     ans = list(ans)
@@ -118,7 +119,9 @@ 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 = filter(
+        lambda w: ".".join(w.torch_version.split(".")[:3]) == torch_version, wheels
+    )
     wheels = list(wheels)
     wheels.sort(reverse=True, key=sort_by_wheel)
     for w in wheels: