fix neuronx key

Signed-off-by: sirutBuasai <sirutbuasai27@outlook.com>

Author: sirutBuasai

docs/src/tables/huggingface_pytorch_table.py

@@ -27,8 +27,8 @@
 DISPLAY_NAMES = {
     "huggingface-pytorch-training": "HuggingFace PyTorch Training",
     "huggingface-pytorch-inference": "HuggingFace PyTorch Inference",
-    "huggingface-pytorch-inference-neuronx": "HuggingFace PyTorch Inference (Neuronx)",
-    "huggingface-pytorch-training-neuronx": "HuggingFace PyTorch Training (Neuronx)",
+    "huggingface-pytorch-inference-neuronx": "HuggingFace PyTorch Inference (NeuronX)",
+    "huggingface-pytorch-training-neuronx": "HuggingFace PyTorch Training (NeuronX)",
     "huggingface-pytorch-trcomp-training": "HuggingFace PyTorch Training Compiler",
 }
 COLUMNS_STANDARD = [
@@ -74,7 +74,7 @@ def parse_tag(tag: str) -> dict:
         result["version"] = match.group(1)
         result["transformers"] = match.group(2)
         accel = match.group(3)
-        result["accelerator"] = "Neuron" if accel == "neuronx" else accel.upper()
+        result["accelerator"] = "NeuronX" if accel == "neuronx" else accel.upper()
         result["python"] = match.group(4)
         extra = match.group(5) or ""
         if extra.startswith("cu"):
@@ -96,7 +96,7 @@ def generate(yaml_data: dict) -> str:
         if not tags:
             continue
 
-        is_neuron = "neuronx" in repo_key
+        is_neuron = "neuron" in repo_key
         columns = COLUMNS_NEURON if is_neuron else COLUMNS_STANDARD
 
         rows = []

docs/src/tables/neuron_table.py

@@ -23,28 +23,29 @@
     "tensorflow-inference-neuronx",
 ]
 DISPLAY_NAMES = {
-    "pytorch-inference-neuronx": "PyTorch Inference (Neuronx)",
-    "pytorch-training-neuronx": "PyTorch Training (Neuronx)",
-    "tensorflow-inference-neuronx": "TensorFlow Inference (Neuronx)",
+    "pytorch-inference-neuronx": "PyTorch Inference (NeuronX)",
+    "pytorch-training-neuronx": "PyTorch Training (NeuronX)",
+    "tensorflow-inference-neuronx": "TensorFlow Inference (NeuronX)",
 }
 COLUMNS = ["Framework", "Python", "SDK", "Accelerator", "Platform", "Example URL"]
 
 
 def parse_tag(tag: str) -> dict:
     """Parse Neuron tag format: 2.8.0-neuronx-py311-sdk2.26.1-ubuntu22.04"""
-    result = {"version": "", "sdk": "", "python": ""}
+    result = {"version": "", "sdk": "", "python": "", "accelerator": ""}
 
     match = re.match(
         r"^(\d+\.\d+\.\d+)-"  # framework version
-        r"neuronx-"  # neuronx marker
+        r"(neuronx)-"  # neuronx marker
         r"(py\d+)-"  # python
         r"sdk([\d.]+)-",  # sdk version
         tag,
     )
     if match:
         result["version"] = match.group(1)
-        result["python"] = match.group(2)
-        result["sdk"] = match.group(3)
+        result["accelerator"] = "NeuronX" if match.group(2) == "neuronx" else "-"
+        result["python"] = match.group(3)
+        result["sdk"] = match.group(4)
 
     return result
 
@@ -75,7 +76,7 @@ def generate(yaml_data: dict) -> str:
                     f"{framework_name} {parsed['version']}",
                     parsed["python"],
                     parsed["sdk"],
-                    "Neuron",
+                    parsed["accelerator"],
                     "SageMaker",
                     build_ecr_url(repo_key, tag),
                 ]