mistral-model-readiness-checker.py

import os
import json
import yaml
from transformers import AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import HfApi

def load_json_config(file_path):
    """Safely load JSON configuration files"""
    try:
        with open(file_path, 'r') as f:
            return json.load(f)
    except Exception as e:
        print(f"❌ Error loading {file_path}: {e}")
        return None

def review_adapter_config(config_path):
    """
    Comprehensive review of adapter configuration
    """
    print("\n🔍 Adapter Configuration Analysis:")
    config = load_json_config(config_path)
    
    if not config:
        return
    
    # Key checks for adapter configuration
    checks = [
        ("Adapter Type", config.get('adapter_type')),
        ("Base Model", config.get('base_model_name_or_path')),
        ("Reduction Factor", config.get('reduction_factor')),
        ("Target Modules", config.get('target_modules'))
    ]
    
    for label, value in checks:
        status = "✅" if value is not None else "❗"
        print(f"{status} {label}: {value}")
    
    # Additional insights
    if 'peft_type' in config:
        print(f"✅ PEFT Type: {config['peft_type']}")

def review_tokenizer_config(config_path):
    """
    Comprehensive review of tokenizer configuration
    """
    print("\n🔍 Tokenizer Configuration Analysis:")
    config = load_json_config(config_path)
    
    if not config:
        return
    
    # Tokenizer key checks
    tokenizer_checks = [
        ("Vocabulary Size", config.get('vocab_size')),
        ("Padding Side", config.get('padding_side')),
        ("Truncation Side", config.get('truncation_side')),
        ("Model Max Length", config.get('model_max_length')),
        ("Special Tokens", config.get('special_tokens_map_file'))
    ]
    
    for label, value in tokenizer_checks:
        status = "✅" if value is not None else "❗"
        print(f"{status} {label}: {value}")

def check_inference_yaml(yaml_path):
    """
    Review inference configuration
    """
    print("\n🚀 Inference Configuration Analysis:")
    try:
        with open(yaml_path, 'r') as f:
            inference_config = yaml.safe_load(f)
        
        # Check key inference parameters
        print("Inference Configuration Details:")
        print(json.dumps(inference_config, indent=2))
    except Exception as e:
        print(f"❌ Error reading inference YAML: {e}")

def model_loading_test(model_id):
    """
    Test model loading and basic generation
    """
    print("\n🧪 Model Loading and Generation Test:")
    try:
        # Load tokenizer
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        print("✅ Tokenizer Loaded Successfully")
        
        # Load model
        model = AutoModelForCausalLM.from_pretrained(model_id)
        print("✅ Model Loaded Successfully")
        
        # Basic generation test
        test_prompt = "Explain machine learning in simple terms:"
        input_ids = tokenizer.encode(test_prompt, return_tensors="pt")
        
        # Generate response
        output = model.generate(input_ids, max_length=100, num_return_sequences=1)
        generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
        
        print("✅ Basic Generation Test Passed")
        print("\n📝 Generated Sample:")
        print(generated_text)
    except Exception as e:
        print(f"❌ Model Loading/Generation Failed: {e}")

def optimize_repository_structure(model_id):
    """
    Provide recommendations for repository optimization
    """
    print("\n🛠️ Repository Optimization Recommendations:")
    
    # Checkpoint management
    print("Checkpoint Management:")
    print("1. Consider keeping only the latest checkpoint (36)")
    print("2. Archive or remove older checkpoints to reduce repository size")
    
    # Configuration file cleanup
    print("\nConfiguration File Optimization:")
    print("1. Ensure consistent naming across checkpoint configurations")
    print("2. Verify that the latest checkpoint has the most up-to-date configs")
    
    # Metadata and documentation
    print("\nDocumentation Recommendations:")
    print("1. Update README.md with:")
    print("   - Model architecture details")
    print("   - Training methodology")
    print("   - Intended use cases")
    print("   - Performance metrics")

def main():
    model_id = "mycholpath/RA-Mistral-7B"
    base_path = "."  # Adjust if needed
    
    # Configuration file paths
    adapter_config_path = os.path.join(base_path, "adapter_config.json")
    tokenizer_config_path = os.path.join(base_path, "tokenizer_config.json")
    inference_yaml_path = os.path.join(base_path, "inference.yaml")
    
    # Run comprehensive diagnostics
    review_adapter_config(adapter_config_path)
    review_tokenizer_config(tokenizer_config_path)
    check_inference_yaml(inference_yaml_path)
    model_loading_test(model_id)
    optimize_repository_structure(model_id)

if __name__ == "__main__":
    main()