create_env.py

import os
import subprocess
import sys

def setup_dependencies():
    os.environ["OMP_NUM_THREADS"] = "4"
    try:
        if os.path.exists('/tmp/deps_installed'):
            return
            
        print("Installing transformers dev version...")
        subprocess.check_call([
            sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-cache-dir",
            "transformers==4.56.0"
        ])
        
        # Fix ml_dtypes / onnx compatibility
        print("Fixing ml_dtypes and onnx compatibility...")
        subprocess.check_call([
            sys.executable, "-m", "pip", "install", "--upgrade",
            "ml_dtypes>=0.5.0", "onnx>=1.17.0"
        ])
        
        # Restore nemo_toolkit compatible versions
        print("Restoring numpy and fsspec to nemo_toolkit compatible versions...")
        subprocess.check_call([
            sys.executable, "-m", "pip", "install", "--force-reinstall",
            "numpy==1.26.4", "fsspec==2024.12.0"
        ])
        
        with open('/tmp/deps_installed', 'w') as f:
            f.write('done')
            
    except Exception as e:
        print(f"Dependencies setup error: {e}")