model_size_checker.py

import os
import sys
from pathlib import Path
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def get_directory_size(directory_path: str) -> float:
    
    total_size = 0
    try:
        for dirpath, dirnames, filenames in os.walk(directory_path):
            for filename in filenames:
                filepath = os.path.join(dirpath, filename)
                if os.path.exists(filepath):
                    total_size += os.path.getsize(filepath)
    except Exception as e:
        logger.warning(f"Error calculating size for {directory_path}: {e}")
    
    return total_size / (1024 * 1024)

def check_model_sizes():
    
    model_dirs = [
        "~/.cache/torch",
        "~/.cache/huggingface",
        ".cache",
        "~/.cache/sentence_transformers"
    ]
    
    total_size = 0
    print("Model Size Analysis:")
    print("=" * 50)
    
    for model_dir in model_dirs:
        expanded_path = os.path.expanduser(model_dir)
        if os.path.exists(expanded_path):
            size_mb = get_directory_size(expanded_path)
            total_size += size_mb
            print(f"{model_dir}: {size_mb:.1f} MB")
        else:
            print(f"{model_dir}: Not found")
    
    print("=" * 50)
    print(f"Total size: {total_size:.1f} MB ({total_size/1024:.2f} GB)")
    
    if total_size > 1024:
        print("❌ WARNING: Total model size exceeds 1GB limit!")
        return False
    else:
        print("✅ Total model size is under 1GB limit")
        return True

def estimate_model_sizes():
    
    print("\nModel Size Estimates:")
    print("=" * 50)
    
    models = {
        "paraphrase-MiniLM-L3-v2": "~61 MB",
        "all-MiniLM-L3-v2": "~61 MB", 
        "all-MiniLM-L6-v2": "~90 MB",
        "distilbart-cnn-12-6": "~1.2 GB"
    }
    
    for model, size in models.items():
        print(f"{model}: {size}")
    
    print("\nCurrent system uses:")
    print("- all-MiniLM-L6-v2 (~90 MB) - UPGRADED for better quality!")
    print("- No large summarization models")
    print("- Total estimated size: ~130-180 MB")

if __name__ == "__main__":
    print("Model Size Checker for Document Intelligence System")
    print("=" * 60)
    
    estimate_model_sizes()
    print()
    
    if check_model_sizes():
        print("\n✅ System is ready to run within 1GB constraint")
    else:
        print("\n❌ System exceeds 1GB constraint - cleanup needed")
        print("Consider running: rm -rf ~/.cache/huggingface ~/.cache/torch")