restore

scripts/train.py

@@ -11,70 +11,8 @@ from datasets import load_dataset
 from myolmoe import MyOlmoeForCausalLM, OlmoeConfig
 import os
 
-import os
-import tempfile
-import datasets
-from datasets import load_dataset
-
-# Create a completely isolated cache directory
-temp_dir = tempfile.mkdtemp()
-print(f"Using temporary cache directory: {temp_dir}")
-
-# Set environment variable to override cache location
-os.environ['HF_DATASETS_CACHE'] = temp_dir
-
-# Also set HF_HOME to ensure clean state
-os.environ['HF_HOME'] = temp_dir
-
-# Disable caching entirely
-datasets.disable_caching()
-
-# 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")
+    print("Starting my COOL OLMoE training script for small experts")
     # Load config - first try from local file, then from pretrained
     config_path = os.path.join("myolmoe", "config.json")
     if os.path.exists(config_path):