hello_metal.py

#!/usr/bin/env python3
"""
Run GPT-OSS using the Metal implementation for Apple Silicon
This requires downloading and converting the model first.
"""

import os
import subprocess
import sys

def setup_metal_model():
    """Download and convert model for Metal usage"""
    
    # Check if model directory exists
    model_dir = "gpt-oss-20b"
    converted_model = "gpt-oss-20b-metal.bin"
    
    if not os.path.exists(converted_model):
        print("Setting up Metal model...")
        
        # Download the model if not present
        if not os.path.exists(model_dir):
            print(f"Downloading {model_dir}...")
            subprocess.run([
                "huggingface-cli", "download", 
                "openai/gpt-oss-20b", 
                "--include", "original/*", 
                "--local-dir", f"{model_dir}/"
            ], check=True)
        
        # Convert the model for Metal
        print(f"Converting model to Metal format...")
        subprocess.run([
            sys.executable, 
            "gpt-oss/gpt_oss/metal/scripts/create-local-model.py",
            "-s", f"{model_dir}/original/",
            "-d", converted_model
        ], check=True)
        
        print(f"Model converted to {converted_model}")
    
    return converted_model

def run_metal_inference(model_path):
    """Run inference using the Metal backend"""
    
    # Import the metal module
    from gpt_oss.generate import main
    import argparse
    
    # Create args for the generate function
    args = argparse.Namespace(
        checkpoint=model_path,
        backend='metal',
        prompt="Explain quantum mechanics clearly and concisely.",
        max_tokens=256,
        interactive=False
    )
    
    # Run generation
    main(args)

if __name__ == "__main__":
    # First ensure gpt-oss is installed with metal support
    print("Installing gpt-oss with Metal support...")
    subprocess.run([
        sys.executable, "-m", "pip", "install", "-e", "./gpt-oss[metal]"
    ], check=False)
    
    # Setup and run
    model_path = setup_metal_model()
    run_metal_inference(model_path)

openai 20b model