Fix semantic evaluator to recover metadata by prediction id

scripts/evaluate_semantic_o1_a1.py

@@ -17,6 +17,8 @@ from collections import defaultdict
 from pathlib import Path
 from typing import Any, Dict, Iterable, List, Optional, Tuple
 
+from datasets import load_dataset
+
 from tmf921_train.utils import aggregate_metrics, parse_json, write_json
 
 
@@ -69,16 +71,31 @@ def any_key_contains(obj: Any, fragments: List[str]) -> bool:
     return False
 
 
+def safe_int(x: Any) -> Optional[int]:
+    try:
+        if x is None:
+            return None
+        return int(x)
+    except Exception:
+        return None
+
+
+def check_contains_if_present(obj: Any, target: Any, *, rel: float = 0.02) -> bool:
+    if target is None:
+        return False
+    return contains_value(obj, target, rel=rel)
+
+
 def semantic_score_o1(example: Dict[str, Any], pred_obj: Any) -> Dict[str, Any]:
     checks = {}
-    checks["sem_sst"] = contains_value(pred_obj, int(example.get("sst")))
-    checks["sem_sd"] = contains_value(pred_obj, example.get("sd"))
+    checks["sem_sst"] = check_contains_if_present(pred_obj, safe_int(example.get("sst")))
+    checks["sem_sd"] = check_contains_if_present(pred_obj, example.get("sd"))
     checks["sem_snssai"] = checks["sem_sst"] and checks["sem_sd"]
-    checks["sem_slice_type"] = contains_value(pred_obj, example.get("slice_type"))
-    checks["sem_latency"] = contains_value(pred_obj, example.get("latency_ms"), rel=0.05)
-    checks["sem_dl"] = contains_value(pred_obj, example.get("dl_throughput_mbps"), rel=0.05)
-    checks["sem_ul"] = contains_value(pred_obj, example.get("ul_throughput_mbps"), rel=0.05)
-    checks["sem_max_ues"] = contains_value(pred_obj, example.get("max_ues"), rel=0.05)
+    checks["sem_slice_type"] = check_contains_if_present(pred_obj, example.get("slice_type"))
+    checks["sem_latency"] = check_contains_if_present(pred_obj, example.get("latency_ms"), rel=0.05)
+    checks["sem_dl"] = check_contains_if_present(pred_obj, example.get("dl_throughput_mbps"), rel=0.05)
+    checks["sem_ul"] = check_contains_if_present(pred_obj, example.get("ul_throughput_mbps"), rel=0.05)
+    checks["sem_max_ues"] = check_contains_if_present(pred_obj, example.get("max_ues"), rel=0.05)
     checks["sem_managed_element_structure"] = any_key_contains(pred_obj, ["ManagedElement", "GNBDU", "NRCell", "cell", "rrmPolicy"])
     checks["sem_rrm_policy_structure"] = any_key_contains(pred_obj, ["rrmPolicy", "prb", "ratio", "quota"])
     checks["sem_cell_parameter_structure"] = any_key_contains(pred_obj, ["arfcn", "pci", "tac", "bandwidth", "bSChannelBw", "cellLocalId"])
@@ -92,18 +109,18 @@ def semantic_score_o1(example: Dict[str, Any], pred_obj: Any) -> Dict[str, Any]:
 
 def semantic_score_a1(example: Dict[str, Any], pred_obj: Any) -> Dict[str, Any]:
     checks = {}
-    checks["sem_sst"] = contains_value(pred_obj, int(example.get("sst")))
-    checks["sem_sd"] = contains_value(pred_obj, example.get("sd"))
+    checks["sem_sst"] = check_contains_if_present(pred_obj, safe_int(example.get("sst")))
+    checks["sem_sd"] = check_contains_if_present(pred_obj, example.get("sd"))
     checks["sem_snssai"] = checks["sem_sst"] and checks["sem_sd"]
-    checks["sem_slice_type"] = contains_value(pred_obj, example.get("slice_type"))
+    checks["sem_slice_type"] = check_contains_if_present(pred_obj, example.get("slice_type"))
     checks["sem_policy_structure"] = any_key_contains(pred_obj, ["policy", "a1", "scope", "objective", "qos", "qoe"])
     checks["sem_prb_or_quota_structure"] = any_key_contains(pred_obj, ["prb", "quota", "resource", "allocation", "ratio"])
     checks["sem_scheduler_structure"] = any_key_contains(pred_obj, ["scheduler", "weight", "priority"])
     checks["sem_5qi_or_qos_structure"] = any_key_contains(pred_obj, ["5qi", "fiveqi", "qci", "qos", "pdb", "per", "gfbr", "mfbr"])
-    checks["sem_latency"] = contains_value(pred_obj, example.get("latency_ms"), rel=0.05)
-    checks["sem_dl"] = contains_value(pred_obj, example.get("dl_throughput_mbps"), rel=0.05)
-    checks["sem_ul"] = contains_value(pred_obj, example.get("ul_throughput_mbps"), rel=0.05)
-    checks["sem_max_ues"] = contains_value(pred_obj, example.get("max_ues"), rel=0.05)
+    checks["sem_latency"] = check_contains_if_present(pred_obj, example.get("latency_ms"), rel=0.05)
+    checks["sem_dl"] = check_contains_if_present(pred_obj, example.get("dl_throughput_mbps"), rel=0.05)
+    checks["sem_ul"] = check_contains_if_present(pred_obj, example.get("ul_throughput_mbps"), rel=0.05)
+    checks["sem_max_ues"] = check_contains_if_present(pred_obj, example.get("max_ues"), rel=0.05)
     core_keys = ["sem_sst", "sem_sd", "sem_snssai", "sem_policy_structure", "sem_prb_or_quota_structure", "sem_5qi_or_qos_structure"]
     kpi_keys = ["sem_latency", "sem_dl", "sem_ul", "sem_max_ues", "sem_scheduler_structure"]
     checks["sem_core_score"] = sum(bool(checks[k]) for k in core_keys) / len(core_keys)
@@ -154,6 +171,7 @@ def summarize(rows: List[Dict[str, Any]]) -> Dict[str, Any]:
 def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--eval_dir", required=True, help="Eval dir with split/predictions.json and optionally normalized_predictions_scored.json")
+    ap.add_argument("--dataset", default="nraptisss/TMF921-intent-to-config-research-sota", help="Dataset used to recover metadata such as sst/sd/KPIs by row id")
     ap.add_argument("--splits", nargs="*", default=None)
     args = ap.parse_args()
     eval_dir = Path(args.eval_dir)
@@ -161,6 +179,15 @@ def main():
         splits = args.splits
     else:
         splits = sorted([p.name for p in eval_dir.iterdir() if p.is_dir() and (p / "predictions.json").exists()])
+    print(f"Loading metadata from {args.dataset} for splits: {splits}")
+    ds = load_dataset(args.dataset)
+    meta_by_split_id = {}
+    needed_cols = ["id", "target_layer", "slice_type", "sst", "sd", "latency_ms", "dl_throughput_mbps", "ul_throughput_mbps", "max_ues"]
+    for split in splits:
+        if split in ds:
+            meta_by_split_id[split] = {str(r["id"]): r for r in ds[split].select_columns([c for c in needed_cols if c in ds[split].column_names])}
+        else:
+            meta_by_split_id[split] = {}
     all_rows = []
     for split in splits:
         pred_path = eval_dir / split / "predictions.json"
@@ -172,6 +199,9 @@ def main():
         scored = []
         for r in rows:
             merged = dict(r)
+            for k, v in meta_by_split_id.get(split, {}).get(str(r.get("id")), {}).items():
+                if k not in merged or merged.get(k) is None:
+                    merged[k] = v
             for k, v in norm.get(r.get("id"), {}).items():
                 if k not in merged:
                     merged[k] = v
@@ -191,5 +221,4 @@ def main():
 
 
 if __name__ == "__main__":
-    import argparse
     main()