Upload 2 files

code/train.py

@@ -134,7 +134,7 @@ class ShardPool:
         except Exception as e:
             logger.warning(f"Failed to list dataset repo {self.repo_id}, trying as model repo. Error: {e}")
             files = self.api.list_repo_files(repo_id=self.repo_id, repo_type="model", token=self.token)
-        self.available_files = [f for f in files if f.startswith("shards/hi_indiccorp_") and f.endswith(".bin")]
+        self.available_files = [f for f in files if f.startswith("shards/") and f.endswith(".bin") and not f.endswith("_val.bin")]
         np.random.shuffle(self.available_files)
         
     def fill(self, ctx_len):
@@ -244,7 +244,7 @@ def main():
     start_step = load_checkpoint(unwrapped_model, optimizer, scaler, api, token, "ViuAI/ViuRec", strict=False)
 
     work_dir = Path("/kaggle/working" if Path("/kaggle/working").exists() else ".")
-    shard_pool = ShardPool(api, "ViuAI/ViuCorp", work_dir / "shards", token, vocab_size)
+    shard_pool = ShardPool(api, "ViuAI/viuai-500m-data", work_dir / "shards", token, vocab_size)
     shard_pool.fill(args.ctx_len)
 
     # Val Data Safely (cap at 20M tokens)
@@ -252,12 +252,12 @@ def main():
     if is_main:
         logger.info("Loading validation data...")
         try:
-            val_files = [f for f in api.list_repo_files(repo_id="ViuAI/ViuCorp", repo_type="dataset", token=token) if f.startswith("val_shards/")]
+            val_files = [f for f in api.list_repo_files(repo_id="ViuAI/viuai-500m-data", repo_type="dataset", token=token) if f.startswith("shards/") and f.endswith("_val.bin")]
         except Exception:
-            val_files = [f for f in api.list_repo_files(repo_id="ViuAI/ViuCorp", repo_type="model", token=token) if f.startswith("val_shards/")]
+            val_files = [f for f in api.list_repo_files(repo_id="ViuAI/viuai-500m-data", repo_type="model", token=token) if f.startswith("shards/") and f.endswith("_val.bin")]
         try:
             if val_files:
-                val_p = hf_hub_download(repo_id="ViuAI/ViuCorp", repo_type="dataset", filename=val_files[0], local_dir=work_dir, token=token)
+                val_p = hf_hub_download(repo_id="ViuAI/viuai-500m-data", repo_type="dataset", filename=val_files[0], local_dir=work_dir, token=token)
                 dtype = np.uint16 if vocab_size <= 65535 else np.uint32
                 val_arr = np.memmap(val_p, dtype=dtype, mode="r")
                 max_toks = 20_000_000