Update src/dataset/atlas_dataset.py: precomputed map tokens support + NCCL timeout

src/dataset/atlas_dataset.py

@@ -103,6 +103,7 @@ class AtlasDataset(Dataset):
         image_path_remap: Optional[str] = None,
         precomputed_det_tokens: Optional[str] = None,
         require_precomputed_det_tokens: bool = False,
+        precomputed_map_tokens: Optional[str] = None,
     ):
         self.json_file = json_file
         self.image_root = image_root
@@ -112,6 +113,7 @@ class AtlasDataset(Dataset):
             self.image_path_remap = (old, new)
         self.precomputed_det_dir = precomputed_det_tokens
         self.require_precomputed = require_precomputed_det_tokens
+        self.precomputed_map_dir = precomputed_map_tokens
         self._precomputed_path_map: Optional[Dict[str, str]] = None
         if self.precomputed_det_dir and os.path.isdir(self.precomputed_det_dir):
             self._precomputed_path_map = {}
@@ -120,6 +122,14 @@ class AtlasDataset(Dataset):
                     if fname.endswith(".pt"):
                         self._precomputed_path_map[fname[:-3]] = os.path.join(root_dir, fname)
             print(f"Precomputed det tokens index: {len(self._precomputed_path_map)} files")
+        self._precomputed_map_path_map: Optional[Dict[str, str]] = None
+        if self.precomputed_map_dir and os.path.isdir(self.precomputed_map_dir):
+            self._precomputed_map_path_map = {}
+            for root_dir, _, files in os.walk(self.precomputed_map_dir):
+                for fname in files:
+                    if fname.endswith(".pt"):
+                        self._precomputed_map_path_map[fname[:-3]] = os.path.join(root_dir, fname)
+            print(f"Precomputed map tokens index: {len(self._precomputed_map_path_map)} files")
         self.tokenizer = tokenizer
         self.max_length = max_length
         self.is_training = is_training
@@ -406,6 +416,11 @@ class AtlasDataset(Dataset):
                 result["precomputed_det"] = pt["detection"]
                 result["precomputed_det_ref"] = pt["detection_ref_points"]
 
+        if self.precomputed_map_dir:
+            mpt = self._load_precomputed_map(item)
+            if mpt is not None:
+                result["precomputed_map"] = mpt
+
         if os.getenv("ATLAS_AUDIT", "0") not in ("", "0", "false", "False"):
             max_samples = int(os.getenv("ATLAS_AUDIT_MAX_SAMPLES", "1"))
             if idx < max_samples:
@@ -636,6 +651,18 @@ class AtlasDataset(Dataset):
         except Exception:
             return None
 
+    def _load_precomputed_map(self, item: Dict) -> Optional[Dict]:
+        if not self.precomputed_map_dir or self._precomputed_map_path_map is None:
+            return None
+        item_id = str(item.get("id", ""))
+        pt_path = self._precomputed_map_path_map.get(item_id)
+        if pt_path is None:
+            return None
+        try:
+            return torch.load(pt_path, map_location="cpu")
+        except Exception:
+            return None
+
     def _get_scene_id(self, item: Dict) -> str:
         if "segment_id" in item and item["segment_id"]:
             return str(item["segment_id"])
@@ -1102,6 +1129,13 @@ def atlas_collate_fn(
         result["precomputed_det"] = torch.stack([item["precomputed_det"] for item in batch])
         result["precomputed_det_ref"] = torch.stack([item["precomputed_det_ref"] for item in batch])
 
+    if all("precomputed_map" in item for item in batch):
+        map_keys = list(batch[0]["precomputed_map"].keys())
+        result["precomputed_map"] = {
+            k: torch.stack([item["precomputed_map"][k] for item in batch])
+            for k in map_keys
+        }
+
     audit_keys = [
         "audit_prompt_len",
         "audit_answer_len",