Upload 2 files

code/train.py

@@ -32,12 +32,14 @@ logger = logging.getLogger(__name__)
 def get_hf_token():
     token = os.environ.get("HF_TOKEN")
     if token:
+        token = token.strip()
         logger.info("Loaded HF_TOKEN from environment.")
         return token
     if HAS_KAGGLE_SECRETS:
         try:
             token = UserSecretsClient().get_secret("HF_TOKEN")
             if token:
+                token = token.strip()
                 os.environ["HF_TOKEN"] = token
                 logger.info("Loaded HF_TOKEN from kaggle_secrets.")
                 return token
@@ -127,7 +129,11 @@ class ShardPool:
         self._refresh_file_list()
         
     def _refresh_file_list(self):
-        files = self.api.list_repo_files(repo_id=self.repo_id, repo_type="dataset", token=self.token)
+        try:
+            files = self.api.list_repo_files(repo_id=self.repo_id, repo_type="dataset", token=self.token)
+        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")]
         np.random.shuffle(self.available_files)
         
@@ -247,6 +253,9 @@ def 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/")]
+        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/")]
+        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)
                 dtype = np.uint16 if vocab_size <= 65535 else np.uint32