import zipfile
import os

def prepare_colab_bundle():
    bundle_name = "ora_colab_training.zip"
    files_to_include = [
        ("important/curated_data/final_ora_dataset.jsonl", "important/curated_data/final_ora_dataset.jsonl"),
        ("scripts/train_ora.py", "scripts/train_ora.py"),
        ("training/requirements.txt", "requirements.txt")
    ]
    
    print(f"Preparing bundle: {bundle_name}...")
    
    with zipfile.ZipFile(bundle_name, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for src, dest in files_to_include:
            if os.path.exists(src):
                print(f"Adding {src} as {dest}...")
                zipf.write(src, arcname=dest)
            else:
                print(f"Warning: {src} not found!")
                
    print(f"\nBundle ready: {os.path.abspath(bundle_name)}")
    print("Next steps: Upload this ZIP to Google Colab, unzip it, and use the unsloth version of the training script.")

if __name__ == "__main__":
    prepare_colab_bundle()