Update script.py

script.py

@@ -4,8 +4,8 @@ os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface/transformers"
 os.environ["HF_HUB_CACHE"] = "/tmp/huggingface/hub"
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
-#HF_DATASET = os.environ.get("HF_DATASET")
-HF_DATASET = "DevWild/autotrain-pr0b0rk"
+HF_DATASET = os.environ.get("HF_DATASET")
+#HF_DATASET = "DevWild/autotrain-pr0b0rk"
 repo_id = os.environ.get("MODEL_REPO_ID")
 
 from huggingface_hub import snapshot_download, delete_repo, metadata_update
@@ -126,7 +126,7 @@ def run_training(hf_dataset_path: str):
         commands = "git clone https://github.com/DevW1ld/ai-toolkit.git ai-toolkit && cd ai-toolkit && git submodule update --init --recursive"
         subprocess.run(commands, shell=True)
 
-    patch_ai_toolkit_typing()
+    # patch_ai_toolkit_typing()
     commands = f"python run.py {os.path.join(dataset_dir, 'config.yaml')}"
     process = subprocess.Popen(commands, shell=True, cwd="ai-toolkit", env=os.environ, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,)
 
@@ -137,19 +137,19 @@ def run_training(hf_dataset_path: str):
     
     return process, dataset_dir
 
-def patch_ai_toolkit_typing():
-    config_path = "ai-toolkit/toolkit/config_modules.py"
-    if os.path.exists(config_path):
-        with open(config_path, "r") as f:
-            content = f.read()
+#def patch_ai_toolkit_typing():
+#    config_path = "ai-toolkit/toolkit/config_modules.py"
+#    if os.path.exists(config_path):
+#        with open(config_path, "r") as f:
+#            content = f.read()
 
-        content = content.replace("torch.Tensor | None", "Optional[torch.Tensor]")
+#        content = content.replace("torch.Tensor | None", "Optional[torch.Tensor]")
 
-        with open(config_path, "w") as f:
-            f.write(content)
-        print("✅ Patched ai-toolkit typing for torch.Tensor | None → Optional[torch.Tensor]")
-    else:
-        print("⚠️ Could not patch config_modules.py — file not found")
+#        with open(config_path, "w") as f:
+#            f.write(content)
+#        print("✅ Patched ai-toolkit typing for torch.Tensor | None → Optional[torch.Tensor]")
+#    else:
+#        print("⚠️ Could not patch config_modules.py — file not found")
 
 
 if __name__ == "__main__":