models.py

"""
Module containing model recommendations and configurations for the Qwen2.5 VL application.
"""

# Dictionary of recommended models with their specifications
RECOMMENDED_MODELS = {
    "Qwen2.5-VL-7B-Instruct": {
        "id": "Qwen/Qwen2.5-VL-7B-Instruct",
        "description": "7B parameter vision-language model with instruction tuning",
        "dtype": "bfloat16",
        "device": "cuda"
    },
    "Qwen2.5-VL-3B-Instruct": {
        "id": "Qwen/Qwen2.5-VL-3B-Instruct",
        "description": "3B parameter vision-language model with instruction tuning",
        "dtype": "bfloat16",
        "device": "cuda"
    },
    "Qwen2-VL-7B": {
        "id": "Qwen/Qwen2-VL-7B",
        "description": "7B parameter vision-language model (Qwen2 series)",
        "dtype": "bfloat16",
        "device": "cuda"
    },
    "Qwen2-VL-7B-Instruct": {
        "id": "Qwen/Qwen2-VL-7B-Instruct",
        "description": "7B parameter vision-language model with instruction tuning (Qwen2 series)",
        "dtype": "bfloat16",
        "device": "cuda"
    }
}

# Default generation parameters
DEFAULT_GENERATION_PARAMS = {
    "max_new_tokens": 1024,
    "temperature": 0.7,
    "top_p": 0.9,
    "top_k": 50,
    "repetition_penalty": 1.0,
    "do_sample": True,
    "num_beams": 1,
    "early_stopping": True,
    "length_penalty": 1.0,
    "no_repeat_ngram_size": 0
}

# Generation parameter presets for quick selection
GENERATION_PRESETS = {
    "Default": DEFAULT_GENERATION_PARAMS,
    "Creative": {
        "max_new_tokens": 1024,
        "temperature": 0.9,
        "top_p": 0.95,
        "top_k": 60,
        "repetition_penalty": 1.1,
        "do_sample": True,
        "num_beams": 1,
        "early_stopping": True,
        "length_penalty": 1.0,
        "no_repeat_ngram_size": 0
    },
    "Precise": {
        "max_new_tokens": 1024,
        "temperature": 0.5,
        "top_p": 0.85,
        "top_k": 40,
        "repetition_penalty": 1.2,
        "do_sample": True,
        "num_beams": 3,
        "early_stopping": True,
        "length_penalty": 1.0,
        "no_repeat_ngram_size": 3
    },
    "Deterministic": {
        "max_new_tokens": 1024,
        "temperature": 0.0,
        "top_p": 1.0,
        "top_k": 50,
        "repetition_penalty": 1.0,
        "do_sample": False,
        "num_beams": 5,
        "early_stopping": True,
        "length_penalty": 1.0,
        "no_repeat_ngram_size": 0
    }
}

def get_model_info(model_name):
    """
    Returns the model information for a given model name.
    
    Args:
        model_name (str): Name of the model
        
    Returns:
        dict: Model specifications
    """
    return RECOMMENDED_MODELS.get(model_name, RECOMMENDED_MODELS["Qwen2.5-VL-7B-Instruct"])

def get_model_list():
    """
    Returns a list of available models for selection.
    
    Returns:
        list: List of model names
    """
    return list(RECOMMENDED_MODELS.keys())

def get_preset_list():
    """
    Returns a list of available parameter presets.
    
    Returns:
        list: List of preset names
    """
    return list(GENERATION_PRESETS.keys())

def get_preset_params(preset_name):
    """
    Returns the generation parameters for a given preset.
    
    Args:
        preset_name (str): Name of the preset
        
    Returns:
        dict: Generation parameters
    """
    return GENERATION_PRESETS.get(preset_name, DEFAULT_GENERATION_PARAMS)

def get_preset_description(preset_name):
    """
    Returns a description for the given preset.
    
    Args:
        preset_name (str): Name of the preset
        
    Returns:
        str: Description of the preset
    """
    descriptions = {
        "Default": "Balanced parameters suitable for most use cases",
        "Creative": "Higher temperature and diversity for more creative outputs",
        "Precise": "Lower temperature with beam search for more precise, focused responses",
        "Deterministic": "Greedy decoding with beam search for deterministic, consistent outputs"
    }
    return descriptions.get(preset_name, "Custom parameters")