claude attempt 2 dataset

scripts/train.py

@@ -11,31 +11,67 @@ from datasets import load_dataset
 from myolmoe import MyOlmoeForCausalLM, OlmoeConfig
 import os
 
+import os
+import tempfile
 import datasets
+from datasets import load_dataset
 
-# Clear the cache
-datasets.disable_caching()
+# Create a completely isolated cache directory
+temp_dir = tempfile.mkdtemp()
+print(f"Using temporary cache directory: {temp_dir}")
 
-# Or manually clear cache
-import shutil
-import os
-cache_dir = os.path.expanduser("~/.cache/huggingface/datasets")
-if os.path.exists(cache_dir):
-    shutil.rmtree(cache_dir)
+# Set environment variable to override cache location
+os.environ['HF_DATASETS_CACHE'] = temp_dir
 
-#
-import datasets
-import os
+# Also set HF_HOME to ensure clean state
+os.environ['HF_HOME'] = temp_dir
 
-# Check current cache directory
-print("Current cache directory:", datasets.config.HF_DATASETS_CACHE)
-print("Cache exists:", os.path.exists(datasets.config.HF_DATASETS_CACHE))
+# Disable caching entirely
+datasets.disable_caching()
 
-# List contents if it exists
-if os.path.exists(datasets.config.HF_DATASETS_CACHE):
-    print("Cache contents:")
-    for item in os.listdir(datasets.config.HF_DATASETS_CACHE):
-        print(f"  {item}")
+# Now try loading
+try:
+    dataset = load_dataset(
+        "allenai/tulu-v2-sft-mixture", 
+        split="train",
+        cache_dir=temp_dir,
+        download_mode="force_redownload"
+    )
+    print(f"SUCCESS! Loaded {len(dataset)} examples")
+except Exception as e:
+    print(f"Still failing: {e}")
+    print("Let's try streaming approach...")
+    
+    # Try streaming instead
+    try:
+        dataset = load_dataset(
+            "allenai/tulu-v2-sft-mixture", 
+            split="train",
+            streaming=True
+        )
+        print("Streaming dataset loaded successfully!")
+        
+        # Convert streaming to regular dataset (sample first 1000 for testing)
+        dataset_list = []
+        for i, example in enumerate(dataset):
+            dataset_list.append(example)
+            if i >= 1000:  # Just for testing
+                break
+        
+        from datasets import Dataset
+        dataset = Dataset.from_list(dataset_list)
+        print(f"Converted to regular dataset with {len(dataset)} examples")
+        
+    except Exception as e2:
+        print(f"Streaming also failed: {e2}")
+        print("Let's try a different dataset temporarily...")
+        
+        # Try a smaller, simpler dataset first
+        try:
+            test_dataset = load_dataset("imdb", split="train[:100]")
+            print("Simple dataset works - issue is specific to tulu-v2-sft-mixture")
+        except Exception as e3:
+            print(f"Even simple dataset fails: {e3}")
 #
 def main():
     print("Starting Training Script for my cool OLMoE")