refactor: improve label functions

Author: Sigmanificient

src/analyzer/data_loader.py

@@ -1,20 +1,23 @@
-from typing import List, Tuple
+from collections import defaultdict
+from typing import DefaultDict, List, Tuple, TypedDict, cast
 
 from chess_utils import fen_to_tensor
 
-from .labels import label_to_vector
+from .labels import label_to_vector, vector_to_label
 
+FenVec = list[float]
 
-def load_chessfile_predict(filepath: str) -> List[str]:
+
+def load_chessfile_predict(filepath: str) -> list[str]:
     with open(filepath, "r", encoding="utf-8") as f:
         return [line.strip() for line in f if line.strip()]
 
 
 def load_chessfile_train(
     filepath: str, encoding: str = "simple"
-) -> List[Tuple[List[float], List[float]]]:
+) -> list[tuple[FenVec, list[float]]]:
 
-    dataset: List[Tuple[List[float], List[float]]] = []
+    dataset: list[tuple[FenVec, list[float]]] = []
 
     with open(filepath, "r", encoding="utf-8") as f:
         for line_num, line in enumerate(f, 1):
@@ -37,3 +40,21 @@ def load_chessfile_train(
                 continue
 
     return dataset
+
+
+class Dataset(TypedDict):
+    nothing: list[FenVec]
+    checkmate_white: list[FenVec]
+    checkmate_black: list[FenVec]
+    check_white: list[FenVec]
+    check_black: list[FenVec]
+
+
+def sort_dataset(ds: list[tuple[FenVec, list[float]]]) -> Dataset:
+    out: DefaultDict[str, list[FenVec]] = defaultdict(list)
+
+    for fen, expected in ds:
+        key = vector_to_label(expected).lower().replace(" ", "_")
+        out[key].append(fen)
+
+    return cast(Dataset, out)

src/analyzer/labels.py

@@ -1,42 +1,30 @@
-from typing import List
-
-
-def label_to_vector(label: str) -> List[float]:
-    label = label.strip()
-
-    if label == "Nothing":
-        return [1.0, 0.0, 0.0, 0.0, 0.0]
-    elif label in ("Check White", "CheckWhite"):
-        return [0.0, 1.0, 0.0, 0.0, 0.0]
-    elif label in ("Check Black", "CheckBlack"):
-        return [0.0, 0.0, 1.0, 0.0, 0.0]
-    elif label in ("Checkmate White", "CheckmateWhite"):
-        return [0.0, 0.0, 0.0, 1.0, 0.0]
-    elif label in ("Checkmate Black", "CheckmateBlack"):
-        return [0.0, 0.0, 0.0, 0.0, 1.0]
-    elif label == "Check":
-        return [0.0, 1.0, 0.0]
-    elif label == "Checkmate":
-        return [0.0, 0.0, 1.0]
-    else:
-        raise ValueError(f"Unknown label: {label}")
-
-
-def vector_to_label(vec: List[float], mode: str = "auto") -> str:
-    if mode == "auto":
-        mode = "detailed" if len(vec) == 5 else "basic"
-
-    pred_idx = max(range(len(vec)), key=lambda i: vec[i])
-
-    if mode == "detailed" and len(vec) == 5:
-        labels = [
-            "Nothing",
-            "Check White",
-            "Check Black",
-            "Checkmate White",
-            "Checkmate Black",
-        ]
-        return labels[pred_idx]
-    else:
-        labels = ["Nothing", "Check", "Checkmate"]
-        return labels[min(pred_idx, 2)]
+from typing import List, Literal
+
+Label = Literal[
+    "Nothing",
+    "Check White",
+    "Check Black",
+    "Checkmate White",
+    "Checkmate Black",
+]
+
+LABELS: list[Label] = [
+    "Nothing",
+    "Check White",
+    "Check Black",
+    "Checkmate White",
+    "Checkmate Black",
+]
+
+
+def label_to_vector(label: Label) -> List[float]:
+    vec = [0.0] * 5
+    vec[LABELS.index(label)] = 1.0
+
+    return vec
+
+
+def vector_to_label(vec: List[float]) -> Label:
+    int_vec = [e > 0.5 for e in vec]
+
+    return LABELS[int_vec.index(True)]