dash QOL

packages/node/meshnet_node/safetensors_selection.py

@@ -118,6 +118,23 @@ def select_files_for_layers_from_index(
     return selected
 
 
+def select_tensor_names_for_layers_from_index(
+    weight_map: dict[str, str],
+    start_layer: int,
+    end_layer: int,
+    *,
+    total_layers: int | None = None,
+) -> set[str]:
+    """Pure variant that returns checkpoint tensor names instead of file paths."""
+    selected: set[str] = set()
+    for tensor_name, rel_file in weight_map.items():
+        if not isinstance(tensor_name, str) or not isinstance(rel_file, str):
+            continue
+        if _tensor_belongs_to_range(tensor_name, start_layer, end_layer, total_layers):
+            selected.add(tensor_name)
+    return selected
+
+
 def _tensor_belongs_to_range(
     tensor_name: str,
     start_layer: int,