anu151105/agentic-browser

1#!/usr/bin/env python32"""3Script to download and cache models for offline use.4"""5import os6import sys7import argparse8import logging9from pathlib import Path10from huggingface_hub import snapshot_download11 12# Set up logging13logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')14logger = logging.getLogger(__name__)15 16# Add the project root to the path17project_root = Path(__file__).parent.parent18sys.path.insert(0, str(project_root))19 20# Try multiple ways to import the config21try:22    from config.model_config import DEFAULT_MODELS23except ImportError as e:24    logger.error(f"Failed to import config: {e}")25    logger.info(f"Current working directory: {os.getcwd()}")26    logger.info(f"Project root: {project_root}")27    logger.info(f"Python path: {sys.path}")28    29    # Try alternative import approach30    try:31        from src.config.model_config import DEFAULT_MODELS32        logger.info("Successfully imported config from src.config.model_config")33    except ImportError as e2:34        logger.error(f"Also failed to import from src.config.model_config: {e2}")35        raise ImportError("Could not import model configuration. Please check your Python path and module structure.")36 37def download_model(model_name: str, cache_dir: str = None):38    """Download a model from Hugging Face Hub if it doesn't exist locally."""39    if model_name not in DEFAULT_MODELS:40        raise ValueError(f"Unknown model: {model_name}")41    42    config = DEFAULT_MODELS[model_name]43    model_path = config.model_path44    45    # Use cache_dir if provided, otherwise use the model's path46    if cache_dir:47        model_path = os.path.join(cache_dir, os.path.basename(model_path))48    49    print(f"Downloading {model_name} to {model_path}...")50    51    # Create model directory if it doesn't exist52    os.makedirs(model_path, exist_ok=True)53    54    # Download the model55    snapshot_download(56        repo_id=config.model_id,57        local_dir=model_path,58        local_dir_use_symlinks=True,59        ignore_patterns=["*.h5", "*.ot", "*.msgpack"],60    )61    62    print(f"Successfully downloaded {model_name} to {model_path}")63 64def main():65    parser = argparse.ArgumentParser(description="Download and cache models for offline use")66    parser.add_argument(67        "--model",68        type=str,69        default="all",70        help="Model to download (default: all)"71    )72    parser.add_argument(73        "--cache-dir",74        type=str,75        default=None,76        help="Directory to cache models (default: model's default path)"77    )78    79    args = parser.parse_args()80    81    if args.model.lower() == "all":82        for model_name in DEFAULT_MODELS.keys():83            try:84                download_model(model_name, args.cache_dir)85            except Exception as e:86                print(f"Error downloading {model_name}: {e}")87    else:88        if args.model not in DEFAULT_MODELS:89            print(f"Error: Unknown model {args.model}")90            print(f"Available models: {', '.join(DEFAULT_MODELS.keys())}")91            return92        download_model(args.model, args.cache_dir)93 94if __name__ == "__main__":95    main()96