nexaai 1.0.9__cp310-cp310-macosx_14_0_universal2.whl → 1.0.11rc1__cp310-cp310-macosx_14_0_universal2.whl

nexaai/utils/manifest_utils.py

@@ -1,280 +0,0 @@
-"""
-Manifest and metadata utilities for handling nexa.manifest files and model metadata.
-
-This module provides utilities to:
-- Load and save nexa.manifest files
-- Create GGUF and MLX manifests
-- Process manifest metadata (handle null fields, fetch avatars, etc.)
-- Manage backward compatibility with old download_metadata.json files
-"""
-
-import os
-import json
-from datetime import datetime
-from typing import Dict, Any, List, Optional
-
-from .quantization_utils import (
-    extract_quantization_from_filename, 
-    detect_quantization_for_mlx
-)
-from .model_types import (
-    PIPELINE_TO_MODEL_TYPE,
-    MODEL_TYPE_TO_PIPELINE
-)
-
-
-def process_manifest_metadata(manifest: Dict[str, Any], repo_id: str) -> Dict[str, Any]:
-    """Process manifest metadata to handle null/missing fields."""
-    # Handle pipeline_tag
-    pipeline_tag = manifest.get('pipeline_tag')
-    if not pipeline_tag:
-        # Reverse map from ModelType if available
-        model_type = manifest.get('ModelType')
-        pipeline_tag = MODEL_TYPE_TO_PIPELINE.get(model_type) if model_type else None
-    
-    # Handle download_time - keep as null if missing
-    download_time = manifest.get('download_time')
-    
-    # Handle avatar_url - fetch on-the-fly if missing/null
-    avatar_url = manifest.get('avatar_url')
-    if not avatar_url:
-        try:
-            from .avatar_fetcher import get_avatar_url_for_repo
-            avatar_url = get_avatar_url_for_repo(repo_id)
-        except Exception:
-            # If fetching fails, leave as None
-            avatar_url = None
-    
-    # Return processed metadata
-    processed_manifest = manifest.copy()
-    processed_manifest.update({
-        'pipeline_tag': pipeline_tag,
-        'download_time': download_time,
-        'avatar_url': avatar_url
-    })
-    
-    return processed_manifest
-
-
-def load_nexa_manifest(directory_path: str) -> Dict[str, Any]:
-    """Load manifest from nexa.manifest if it exists."""
-    manifest_path = os.path.join(directory_path, 'nexa.manifest')
-    if os.path.exists(manifest_path):
-        try:
-            with open(manifest_path, 'r', encoding='utf-8') as f:
-                return json.load(f)
-        except (json.JSONDecodeError, IOError):
-            pass
-    return {}
-
-
-def load_download_metadata(directory_path: str, repo_id: Optional[str] = None) -> Dict[str, Any]:
-    """Load download metadata from nexa.manifest if it exists, fallback to old format."""
-    # First try to load from new manifest format
-    manifest = load_nexa_manifest(directory_path)
-    if manifest and repo_id:
-        # Process the manifest to handle null/missing fields
-        return process_manifest_metadata(manifest, repo_id)
-    elif manifest:
-        # Return manifest as-is if no repo_id provided (for backward compatibility)
-        return manifest
-    
-    # Fallback to old format for backward compatibility
-    old_metadata_path = os.path.join(directory_path, 'download_metadata.json')
-    if os.path.exists(old_metadata_path):
-        try:
-            with open(old_metadata_path, 'r', encoding='utf-8') as f:
-                return json.load(f)
-        except (json.JSONDecodeError, IOError):
-            pass
-    return {}
-
-
-def save_download_metadata(directory_path: str, metadata: Dict[str, Any]) -> None:
-    """Save download metadata to nexa.manifest in the new format."""
-    manifest_path = os.path.join(directory_path, 'nexa.manifest')
-    try:
-        with open(manifest_path, 'w', encoding='utf-8') as f:
-            json.dump(metadata, f, indent=2)
-    except IOError:
-        # If we can't save metadata, don't fail the download
-        pass
-
-
-def create_gguf_manifest(repo_id: str, files: List[str], directory_path: str, old_metadata: Dict[str, Any]) -> Dict[str, Any]:
-    """Create GGUF format manifest."""
-    
-    # Load existing manifest to merge GGUF files if it exists
-    existing_manifest = load_nexa_manifest(directory_path)
-    
-    model_files = {}
-    if existing_manifest and "ModelFile" in existing_manifest:
-        model_files = existing_manifest["ModelFile"].copy()
-    
-    # Process GGUF files
-    for file_name in files:
-        if file_name.endswith('.gguf'):
-            # Use the new enum-based quantization extraction
-            quantization_type = extract_quantization_from_filename(file_name)
-            quant_level = quantization_type.value if quantization_type else "UNKNOWN"
-            
-            file_path = os.path.join(directory_path, file_name)
-            file_size = 0
-            if os.path.exists(file_path):
-                try:
-                    file_size = os.path.getsize(file_path)
-                except (OSError, IOError):
-                    pass
-            
-            model_files[quant_level] = {
-                "Name": file_name,
-                "Downloaded": True,
-                "Size": file_size
-            }
-    
-    manifest = {
-        "Name": repo_id,
-        "ModelType": PIPELINE_TO_MODEL_TYPE.get(old_metadata.get('pipeline_tag'), "other"),
-        "PluginId": "llama_cpp",
-        "ModelFile": model_files,
-        "MMProjFile": {
-            "Name": "",
-            "Downloaded": False,
-            "Size": 0
-        },
-        "TokenizerFile": {
-            "Name": "",
-            "Downloaded": False,
-            "Size": 0
-        },
-        "ExtraFiles": None,
-        # Preserve old metadata fields
-        "pipeline_tag": old_metadata.get('pipeline_tag'),
-        "download_time": old_metadata.get('download_time'),
-        "avatar_url": old_metadata.get('avatar_url')
-    }
-    
-    return manifest
-
-
-def create_mlx_manifest(repo_id: str, files: List[str], directory_path: str, old_metadata: Dict[str, Any]) -> Dict[str, Any]:
-    """Create MLX format manifest."""
-    
-    model_files = {}
-    extra_files = []
-    
-    # Try different methods to extract quantization for MLX models
-    quantization_type = detect_quantization_for_mlx(repo_id, directory_path)
-    
-    # Use the detected quantization or default to "DEFAULT"
-    quant_level = quantization_type.value if quantization_type else "DEFAULT"
-    
-    for file_name in files:
-        file_path = os.path.join(directory_path, file_name)
-        file_size = 0
-        if os.path.exists(file_path):
-            try:
-                file_size = os.path.getsize(file_path)
-            except (OSError, IOError):
-                pass
-        
-        # Check if this is a main model file (safetensors but not index files)
-        if (file_name.endswith('.safetensors') and not file_name.endswith('.index.json')):
-            model_files[quant_level] = {
-                "Name": file_name,
-                "Downloaded": True,
-                "Size": file_size
-            }
-        else:
-            # Add to extra files
-            extra_files.append({
-                "Name": file_name,
-                "Downloaded": True,
-                "Size": file_size
-            })
-    
-    manifest = {
-        "Name": repo_id,
-        "ModelType": PIPELINE_TO_MODEL_TYPE.get(old_metadata.get('pipeline_tag'), "other"),
-        "PluginId": "mlx",
-        "ModelFile": model_files,
-        "MMProjFile": {
-            "Name": "",
-            "Downloaded": False,
-            "Size": 0
-        },
-        "TokenizerFile": {
-            "Name": "",
-            "Downloaded": False,
-            "Size": 0
-        },
-        "ExtraFiles": extra_files if extra_files else None,
-        # Preserve old metadata fields
-        "pipeline_tag": old_metadata.get('pipeline_tag'),
-        "download_time": old_metadata.get('download_time'),
-        "avatar_url": old_metadata.get('avatar_url')
-    }
-    
-    return manifest
-
-
-def detect_model_type(files: List[str]) -> str:
-    """Detect if this is a GGUF or MLX model based on file extensions."""
-    has_gguf = any(f.endswith('.gguf') for f in files)
-    has_safetensors = any(f.endswith('.safetensors') or 'safetensors' in f for f in files)
-    
-    if has_gguf:
-        return "gguf"
-    elif has_safetensors:
-        return "mlx"
-    else:
-        # Default to mlx for other types
-        return "mlx"
-
-
-def create_manifest_from_files(repo_id: str, files: List[str], directory_path: str, old_metadata: Dict[str, Any]) -> Dict[str, Any]:
-    """
-    Create appropriate manifest format based on detected model type.
-    
-    Args:
-        repo_id: Repository ID
-        files: List of files in the model directory
-        directory_path: Path to the model directory
-        old_metadata: Existing metadata (pipeline_tag, download_time, avatar_url)
-        
-    Returns:
-        Dict containing the appropriate manifest format
-    """
-    model_type = detect_model_type(files)
-    
-    if model_type == "gguf":
-        return create_gguf_manifest(repo_id, files, directory_path, old_metadata)
-    else:  # mlx or other
-        return create_mlx_manifest(repo_id, files, directory_path, old_metadata)
-
-
-def save_manifest_with_files_metadata(repo_id: str, local_dir: str, old_metadata: Dict[str, Any]) -> None:
-    """
-    Create and save manifest based on files found in the directory.
-    
-    Args:
-        repo_id: Repository ID
-        local_dir: Local directory containing the model files
-        old_metadata: Existing metadata to preserve
-    """
-    # Get list of files in the directory
-    files = []
-    try:
-        for root, dirs, filenames in os.walk(local_dir):
-            for filename in filenames:
-                # Store relative path from the directory
-                rel_path = os.path.relpath(os.path.join(root, filename), local_dir)
-                files.append(rel_path)
-    except (OSError, IOError):
-        pass
-    
-    # Create appropriate manifest
-    manifest = create_manifest_from_files(repo_id, files, local_dir, old_metadata)
-    
-    # Save manifest
-    save_download_metadata(local_dir, manifest)

nexaai/utils/model_types.py

@@ -1,47 +0,0 @@
-"""
-Model type mappings for HuggingFace pipeline tags to our internal model types.
-
-This module provides centralized model type mapping functionality to avoid
-circular imports between other utility modules.
-"""
-
-from enum import Enum
-from typing import Dict
-
-
-class ModelTypeMapping(Enum):
-    """Enum for mapping HuggingFace pipeline_tag to our ModelType."""
-    TEXT_GENERATION = ("text-generation", "llm")
-    IMAGE_TEXT_TO_TEXT = ("image-text-to-text", "vlm")
-    
-    def __init__(self, pipeline_tag: str, model_type: str):
-        self.pipeline_tag = pipeline_tag
-        self.model_type = model_type
-
-
-# Create mapping dictionaries from the enum
-PIPELINE_TO_MODEL_TYPE: Dict[str, str] = {
-    mapping.pipeline_tag: mapping.model_type 
-    for mapping in ModelTypeMapping
-}
-
-MODEL_TYPE_TO_PIPELINE: Dict[str, str] = {
-    mapping.model_type: mapping.pipeline_tag 
-    for mapping in ModelTypeMapping
-}
-
-
-def map_pipeline_tag_to_model_type(pipeline_tag: str) -> str:
-    """Map HuggingFace pipeline_tag to our ModelType."""
-    if not pipeline_tag:
-        return "other"
-    
-    return PIPELINE_TO_MODEL_TYPE.get(pipeline_tag, "other")
-
-
-def map_model_type_to_pipeline_tag(model_type: str) -> str:
-    """Reverse map ModelType back to HuggingFace pipeline_tag."""
-    if not model_type:
-        return None
-    
-    return MODEL_TYPE_TO_PIPELINE.get(model_type)

nexaai/utils/quantization_utils.py

@@ -1,239 +0,0 @@
-"""
-Quantization utilities for extracting quantization types from model files and configurations.
-
-This module provides utilities to extract quantization information from:
-- GGUF model filenames
-- MLX model repository IDs  
-- MLX model config.json files
-"""
-
-import os
-import json
-import re
-import logging
-from enum import Enum
-from typing import Optional
-
-# Set up logger
-logger = logging.getLogger(__name__)
-
-
-class QuantizationType(str, Enum):
-    """Enum for GGUF and MLX model quantization types."""
-    # GGUF quantization types
-    BF16 = "BF16"
-    F16 = "F16"
-    Q2_K = "Q2_K"
-    Q2_K_L = "Q2_K_L"
-    Q3_K_M = "Q3_K_M"
-    Q3_K_S = "Q3_K_S"
-    Q4_0 = "Q4_0"
-    Q4_1 = "Q4_1"
-    Q4_K_M = "Q4_K_M"
-    Q4_K_S = "Q4_K_S"
-    Q5_K_M = "Q5_K_M"
-    Q5_K_S = "Q5_K_S"
-    Q6_K = "Q6_K"
-    Q8_0 = "Q8_0"
-    MXFP4 = "MXFP4"
-    MXFP8 = "MXFP8"
-    
-    # MLX bit-based quantization types
-    BIT_1 = "1BIT"
-    BIT_2 = "2BIT"
-    BIT_3 = "3BIT"
-    BIT_4 = "4BIT"
-    BIT_5 = "5BIT"
-    BIT_6 = "6BIT"
-    BIT_7 = "7BIT"
-    BIT_8 = "8BIT"
-    BIT_16 = "16BIT"
-
-
-def extract_quantization_from_filename(filename: str) -> Optional[QuantizationType]:
-    """
-    Extract quantization type from filename.
-    
-    Args:
-        filename: The filename to extract quantization from
-        
-    Returns:
-        QuantizationType enum value or None if not found
-    """
-    # Define mapping from lowercase patterns to enum values
-    # Include "." to ensure precise matching (e.g., "q4_0." not "q4_0_xl")
-    pattern_to_enum = {
-        'bf16.': QuantizationType.BF16,
-        'f16.': QuantizationType.F16,  # Add F16 support
-        'q2_k_l.': QuantizationType.Q2_K_L,  # Check Q2_K_L before Q2_K to avoid partial match
-        'q2_k.': QuantizationType.Q2_K,
-        'q3_k_m.': QuantizationType.Q3_K_M,
-        'q3_ks.': QuantizationType.Q3_K_S,
-        'q4_k_m.': QuantizationType.Q4_K_M,
-        'q4_k_s.': QuantizationType.Q4_K_S,
-        'q4_0.': QuantizationType.Q4_0,
-        'q4_1.': QuantizationType.Q4_1,
-        'q5_k_m.': QuantizationType.Q5_K_M,
-        'q5_k_s.': QuantizationType.Q5_K_S,
-        'q6_k.': QuantizationType.Q6_K,
-        'q8_0.': QuantizationType.Q8_0,
-        'mxfp4.': QuantizationType.MXFP4,
-        'mxfp8.': QuantizationType.MXFP8,
-    }
-    
-    filename_lower = filename.lower()
-    
-    # Check longer patterns first to avoid partial matches
-    # Sort by length descending to check q2_k_l before q2_k, q4_k_m before q4_0, etc.
-    for pattern in sorted(pattern_to_enum.keys(), key=len, reverse=True):
-        if pattern in filename_lower:
-            return pattern_to_enum[pattern]
-    
-    return None
-
-
-def extract_quantization_from_repo_id(repo_id: str) -> Optional[QuantizationType]:
-    """
-    Extract quantization type from repo_id for MLX models by looking for bit patterns.
-    
-    Args:
-        repo_id: The repository ID to extract quantization from
-        
-    Returns:
-        QuantizationType enum value or None if not found
-    """
-    # Define mapping from bit numbers to enum values
-    bit_to_enum = {
-        1: QuantizationType.BIT_1,
-        2: QuantizationType.BIT_2,
-        3: QuantizationType.BIT_3,
-        4: QuantizationType.BIT_4,
-        5: QuantizationType.BIT_5,
-        6: QuantizationType.BIT_6,
-        7: QuantizationType.BIT_7,
-        8: QuantizationType.BIT_8,
-        16: QuantizationType.BIT_16,
-    }
-    
-    # First check for patterns like "4bit", "8bit" etc. (case insensitive)
-    pattern = r'(\d+)bit'
-    matches = re.findall(pattern, repo_id.lower())
-    
-    for match in matches:
-        try:
-            bit_number = int(match)
-            if bit_number in bit_to_enum:
-                logger.debug(f"Found {bit_number}bit quantization in repo_id: {repo_id}")
-                return bit_to_enum[bit_number]
-        except ValueError:
-            continue
-    
-    # Also check for patterns like "-q8", "_Q4" etc.
-    q_pattern = r'[-_]q(\d+)'
-    q_matches = re.findall(q_pattern, repo_id.lower())
-    
-    for match in q_matches:
-        try:
-            bit_number = int(match)
-            if bit_number in bit_to_enum:
-                logger.debug(f"Found Q{bit_number} quantization in repo_id: {repo_id}")
-                return bit_to_enum[bit_number]
-        except ValueError:
-            continue
-    
-    return None
-
-
-def extract_quantization_from_mlx_config(mlx_folder_path: str) -> Optional[QuantizationType]:
-    """
-    Extract quantization type from MLX model's config.json file.
-    
-    Args:
-        mlx_folder_path: Path to the MLX model folder
-        
-    Returns:
-        QuantizationType enum value or None if not found
-    """
-    config_path = os.path.join(mlx_folder_path, "config.json")
-    
-    if not os.path.exists(config_path):
-        logger.debug(f"Config file not found: {config_path}")
-        return None
-    
-    try:
-        with open(config_path, 'r', encoding='utf-8') as f:
-            config = json.load(f)
-        
-        # Look for quantization.bits field
-        quantization_config = config.get("quantization", {})
-        if isinstance(quantization_config, dict):
-            bits = quantization_config.get("bits")
-            if isinstance(bits, int):
-                # Define mapping from bit numbers to enum values
-                bit_to_enum = {
-                    1: QuantizationType.BIT_1,
-                    2: QuantizationType.BIT_2,
-                    3: QuantizationType.BIT_3,
-                    4: QuantizationType.BIT_4,
-                    5: QuantizationType.BIT_5,
-                    6: QuantizationType.BIT_6,
-                    7: QuantizationType.BIT_7,
-                    8: QuantizationType.BIT_8,
-                    16: QuantizationType.BIT_16,
-                }
-                
-                if bits in bit_to_enum:
-                    logger.debug(f"Found {bits}bit quantization in config.json: {config_path}")
-                    return bit_to_enum[bits]
-                else:
-                    logger.debug(f"Unsupported quantization bits value: {bits}")
-        
-    except (json.JSONDecodeError, IOError) as e:
-        logger.warning(f"Error reading config.json from {config_path}: {e}")
-    except Exception as e:
-        logger.warning(f"Unexpected error reading config.json from {config_path}: {e}")
-    
-    return None
-
-
-def extract_gguf_quantization(filename: str) -> str:
-    """
-    Extract quantization level from GGUF filename using the enum-based approach.
-    
-    This function provides backward compatibility by returning a string representation
-    of the quantization type.
-    
-    Args:
-        filename: The GGUF filename
-        
-    Returns:
-        String representation of the quantization type or "UNKNOWN" if not found
-    """
-    quantization_type = extract_quantization_from_filename(filename)
-    if quantization_type:
-        return quantization_type.value
-    return "UNKNOWN"
-
-
-def detect_quantization_for_mlx(repo_id: str, directory_path: str) -> Optional[QuantizationType]:
-    """
-    Detect quantization for MLX models using multiple methods in priority order.
-    
-    Args:
-        repo_id: The repository ID
-        directory_path: Path to the model directory
-        
-    Returns:
-        QuantizationType enum value or None if not found
-    """
-    # Method 1: Extract from repo_id
-    quantization_type = extract_quantization_from_repo_id(repo_id)
-    if quantization_type:
-        return quantization_type
-    
-    # Method 2: Extract from config.json if available
-    quantization_type = extract_quantization_from_mlx_config(directory_path)
-    if quantization_type:
-        return quantization_type
-    
-    return None