examples/basic_usage.py

"""
Basic usage examples for DOM KNN models from Hugging Face Hub.

This script demonstrates how to load and use the models for formula assignment.
"""

from transformers import AutoModel
import numpy as np


def example_1_default_model():
    """Load the default best-performing model."""
    print("=" * 60)
    print("Example 1: Loading default model (21T, K=1, Euclidean)")
    print("=" * 60)
    
    model = AutoModel.from_pretrained(
        "SaeedLab/dom-formula-assignment-using-knn",
        trust_remote_code=True
    )
    
    print(f"[OK] Model loaded successfully")
    print(f"  Config: {model.config}")
    print()


def example_2_variant_by_parameters():
    """Load a specific variant using parameters."""
    print("=" * 60)
    print("Example 2: Loading variant by parameters")
    print("=" * 60)
    
    model = AutoModel.from_pretrained(
        "SaeedLab/dom-formula-assignment-using-knn",
        data_source="7T-21T",
        k_neighbors=1,
        metric="euclidean",
        training_version="ver3",
        trust_remote_code=True
    )
    
    print(f"[OK] Model loaded successfully")
    print(f"  Data source: 7T-21T")
    print(f"  K neighbors: 1")
    print(f"  Metric: euclidean")
    print(f"  Training version: ver3")
    print()


def example_3_variant_by_name():
    """Load a specific variant by direct name."""
    print("=" * 60)
    print("Example 3: Loading variant by name")
    print("=" * 60)
    
    model = AutoModel.from_pretrained(
        "SaeedLab/dom-formula-assignment-using-knn",
        variant="knn_21T_k3_manhattan",
        trust_remote_code=True
    )
    
    print(f"[OK] Model loaded: knn_21T_k3_manhattan")
    print(f"  Config: {model.config}")
    print()


def example_4_make_predictions():
    """Make predictions with the model."""
    print("=" * 60)
    print("Example 4: Making predictions")
    print("=" * 60)
    
    # Load model
    model = AutoModel.from_pretrained(
        "SaeedLab/dom-formula-assignment-using-knn",
        trust_remote_code=True
    )
    
    # Create sample data (replace with your actual mass spec features)
    # Features might include: m/z, O/C ratio, H/C ratio, N/C ratio, etc.
    X_sample = np.array([
        [300.1234, 0.5, 1.2, 0.1],  # Sample 1
        [450.6789, 0.6, 1.5, 0.2],  # Sample 2
        [275.5432, 0.4, 1.1, 0.0],  # Sample 3
    ])
    
    print(f"Input shape: {X_sample.shape}")
    print(f"Making predictions...")
    
    # Get predictions
    predictions = model(X_sample)
    
    print(f"[OK] Predictions shape: {predictions.shape}")
    print(f"  First 3 predictions: {predictions[:3]}")
    print()


def example_5_compare_models():
    """Compare predictions from different model variants."""
    print("=" * 60)
    print("Example 5: Comparing model variants")
    print("=" * 60)
    
    # Sample data
    X_sample = np.array([[300.1234, 0.5, 1.2, 0.1]])
    
    variants = [
        ("21T K=1 Euclidean", {"variant": "knn_21T_k1_euclidean"}),
        ("21T K=3 Manhattan", {"variant": "knn_21T_k3_manhattan"}),
        ("7T-21T K=1 Euclidean", {
            "data_source": "7T-21T",
            "k_neighbors": 1,
            "metric": "euclidean",
            "training_version": "ver3"
        }),
    ]
    
    for name, params in variants:
        model = AutoModel.from_pretrained(
            "SaeedLab/dom-formula-assignment-using-knn",
            trust_remote_code=True,
            **params
        )
        pred = model(X_sample)
        print(f"  {name}: {pred[0]}")
    
    print()


if __name__ == "__main__":
    print("\n" + "=" * 60)
    print("DOM KNN Models - Usage Examples")
    print("=" * 60 + "\n")
    
    try:
        example_1_default_model()
        example_2_variant_by_parameters()
        example_3_variant_by_name()
        example_4_make_predictions()
        example_5_compare_models()
        
        print("=" * 60)
        print("[OK] All examples completed successfully!")
        print("=" * 60)
        
    except Exception as e:
        print(f"\n[ERROR] Error: {e}")
        print("\nNote: These examples require the model to be uploaded to")
        print("      Hugging Face Hub at: SaeedLab/dom-formula-assignment-using-knn")
        print("\nFor local testing, replace the repo ID with a local path.")