__init__.py

from typing import Optional

import torch
from loguru import logger
from transformers.configuration_utils import PretrainedConfig

from lorax_server.models.bloom import BLOOMSharded
from lorax_server.models.causal_lm import CausalLM
from lorax_server.models.flash_causal_lm import FlashCausalLM
from lorax_server.models.galactica import GalacticaSharded
from lorax_server.models.model import Model
from lorax_server.models.mpt import MPTSharded
from lorax_server.models.opt import OPTSharded
from lorax_server.models.santacoder import SantaCoder
from lorax_server.models.seq2seq_lm import Seq2SeqLM
from lorax_server.models.t5 import T5Sharded
from lorax_server.utils.sources import get_s3_model_local_dir
from lorax_server.utils.torch_utils import is_bf16_supported

# The flag below controls whether to allow TF32 on matmul. This flag defaults to False
# in PyTorch 1.12 and later.
torch.backends.cuda.matmul.allow_tf32 = True

# The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.
torch.backends.cudnn.allow_tf32 = True

# Disable gradients
torch.set_grad_enabled(False)

__all__ = [
    "Model",
    "BLOOMSharded",
    "CausalLM",
    "FlashCausalLM",
    "GalacticaSharded",
    "Seq2SeqLM",
    "SantaCoder",
    "OPTSharded",
    "T5Sharded",
    "get_model",
]


def get_model(
    model_id: str,
    adapter_id: str,
    revision: Optional[str],
    sharded: bool,
    quantize: Optional[str],
    compile: bool,
    dtype: Optional[str],
    trust_remote_code: bool,
    source: str,
    adapter_source: str,
    merge_adapter_weights: bool,
    embedding_dim: Optional[int] = None,
) -> Model:
    config_dict = None
    if source == "s3":
        # change the model id to be the local path to the folder so
        # we can load the config_dict locally
        logger.info("Using the local files since we are coming from s3")
        model_path = get_s3_model_local_dir(model_id)
        logger.info(f"model_path: {model_path}")
        config_dict, _ = PretrainedConfig.get_config_dict(
            model_path, revision=revision, trust_remote_code=trust_remote_code
        )
        logger.info(f"config_dict: {config_dict}")
        model_id = str(model_path)
    elif source == "hub":
        config_dict, _ = PretrainedConfig.get_config_dict(
            model_id, revision=revision, trust_remote_code=trust_remote_code
        )
    else:
        raise ValueError(f"Unknown source {source}")

    model_type = config_dict["model_type"]
    is_dtype_provided = dtype is not None
    dtype = dtype or config_dict.get("torch_dtype", "float16")

    if dtype in {"float16", "float32"}:
        dtype = torch.float16
    elif dtype == "bfloat16":
        if not is_bf16_supported():
            if is_dtype_provided:
                raise RuntimeError("bfloat16 is not supported on this device, set --dtype float16.")
            logger.warning("bfloat16 is not supported on this device, falling back to float16")
            dtype = torch.float16
        else:
            dtype = torch.bfloat16
    else:
        try:
            dtype = getattr(torch, dtype)
        except AttributeError:
            raise RuntimeError(f"Unknown dtype {dtype}")

    if "facebook/galactica" in model_id:
        return GalacticaSharded(
            model_id,
            revision,
            quantize=quantize,
            compile=compile,
            dtype=dtype,
            dtypetrust_remote_code=trust_remote_code,
        )

    if model_type == "bert":
        from lorax_server.models.flash_bert import FlashBert

        if config_dict["architectures"][0] == "BertForTokenClassification":
            return FlashBert(model_id, revision=revision, dtype=dtype, classifcation_head=True)
        return FlashBert(model_id, revision=revision, dtype=dtype)

    if model_type == "distilbert":
        from lorax_server.models.flash_distilbert import FlashDistilBert

        if config_dict["architectures"][0] == "DistilBertForMaskedLM":
            return FlashDistilBert(model_id, revision=revision, dtype=dtype)

        if config_dict["architectures"][0] == "DistilBertForTokenClassification":
            return FlashDistilBert(model_id, revision=revision, dtype=dtype, classifcation_head=True)

    if model_type == "xlm-roberta":
        from lorax_server.models.flash_roberta import FlashXlmRoberta

        return FlashXlmRoberta(
            model_id,
            adapter_id,
            adapter_source,
            revision=revision,
            dtype=dtype,
            merge_adapter_weights=merge_adapter_weights,
        )

    flash_causal_lm_kwargs = dict(
        quantize=quantize,
        compile=compile,
        dtype=dtype,
        trust_remote_code=trust_remote_code,
        merge_adapter_weights=merge_adapter_weights,
    )

    if model_id.startswith("bigcode/") or model_type == "gpt_bigcode":
        from lorax_server.models.flash_santacoder import FlashSantacoderSharded

        return FlashSantacoderSharded(
            model_id,
            revision,
            **flash_causal_lm_kwargs,
        )

    if model_type == "bloom":
        return BLOOMSharded(
            model_id,
            revision,
            quantize=quantize,
            compile=compile,
            dtype=dtype,
            trust_remote_code=trust_remote_code,
        )

    if model_type == "mpt":
        return MPTSharded(
            model_id,
            revision,
            quantize=quantize,
            compile=compile,
            trust_remote_code=trust_remote_code,
        )

    if model_type == "gpt_neox":
        from lorax_server.models.flash_neox import FlashNeoXSharded

        return FlashNeoXSharded(
            model_id,
            revision,
            quantize=quantize,
            compile=compile,
            dtype=dtype,
            trust_remote_code=trust_remote_code,
        )

    if model_type == "llama":
        from lorax_server.models.flash_llama import FlashLlama

        return FlashLlama(
            model_id,
            adapter_id,
            adapter_source,
            revision,
            **flash_causal_lm_kwargs,
        )

    if model_type == "gpt2":
        from lorax_server.models.flash_gpt2 import FlashGPT2

        return FlashGPT2(
            model_id,
            adapter_id,
            adapter_source,
            revision,
            **flash_causal_lm_kwargs,
        )

    if model_type in ["RefinedWeb", "RefinedWebModel", "falcon"]:
        from lorax_server.models.flash_rw import FlashRWSharded

        return FlashRWSharded(
            model_id,
            revision,
            **flash_causal_lm_kwargs,
        )

    if model_type == "mistral":
        from lorax_server.models.flash_mistral import FlashMistral

        return FlashMistral(
            model_id,
            adapter_id,
            adapter_source,
            revision,
            **flash_causal_lm_kwargs,
        )

    if model_type == "mixtral":
        from lorax_server.models.flash_mixtral import FlashMixtral

        return FlashMixtral(
            model_id,
            adapter_id,
            adapter_source,
            revision,
            **flash_causal_lm_kwargs,
        )

    if model_type == "qwen":
        from lorax_server.models.flash_qwen import FlashQwen

        return FlashQwen(
            model_id,
            adapter_id,
            adapter_source,
            revision,
            **flash_causal_lm_kwargs,
        )

    if model_type == "qwen2":
        from lorax_server.models.flash_qwen2 import FlashQwen2

        return FlashQwen2(
            model_id,
            adapter_id,
            adapter_source,
            revision,
            embedding_dim=embedding_dim,
            **flash_causal_lm_kwargs,
        )

    if model_type in ["phi-msft", "phi"]:
        from lorax_server.models.flash_phi import FlashPhi

        return FlashPhi(
            model_id,
            adapter_id,
            adapter_source,
            revision,
            **flash_causal_lm_kwargs,
        )

    if model_type == "phi3":
        from lorax_server.models.flash_phi3 import FlashPhi3

        return FlashPhi3(
            model_id,
            adapter_id,
            adapter_source,
            revision,
            **flash_causal_lm_kwargs,
        )

    if model_type == "solar":
        from lorax_server.models.flash_solar import FlashSolar

        return FlashSolar(
            model_id,
            adapter_id,
            adapter_source,
            revision,
            **flash_causal_lm_kwargs,
        )

    if model_type == "gemma":
        from lorax_server.models.flash_gemma import FlashGemma

        return FlashGemma(
            model_id,
            adapter_id,
            adapter_source,
            revision,
            **flash_causal_lm_kwargs,
        )

    if model_type == "gemma2":
        from lorax_server.models.flash_gemma2 import FlashGemma2

        return FlashGemma2(
            model_id,
            adapter_id,
            adapter_source,
            revision,
            **flash_causal_lm_kwargs,
        )

    if model_type == "cohere":
        from lorax_server.models.flash_cohere import FlashCohere

        return FlashCohere(
            model_id,
            adapter_id,
            adapter_source,
            revision,
            **flash_causal_lm_kwargs,
        )

    if model_type == "dbrx":
        from lorax_server.models.flash_dbrx import FlashDbrx

        return FlashDbrx(
            model_id,
            adapter_id,
            adapter_source,
            revision,
            **flash_causal_lm_kwargs,
        )

    if model_type == "llava_next" or model_type == "llava":
        from lorax_server.models.custom_modeling.llava_next import LlavaNextForConditionalGeneration
        from lorax_server.models.vlm_causal_lm import VlmCausalLM

        return VlmCausalLM(
            model_class=LlavaNextForConditionalGeneration,
            model_id=model_id,
            adapter_id=adapter_id,
            adapter_source=adapter_source,
            revision=revision,
            **flash_causal_lm_kwargs,
        )

    if model_type == "mllama":
        from lorax_server.models.custom_modeling.mllama import MllamaForConditionalGeneration
        from lorax_server.models.mllama import MllamaCausalLM, MllamaCausalLMBatch

        return MllamaCausalLM(
            model_id=model_id,
            model_class=MllamaForConditionalGeneration,
            batch_class=MllamaCausalLMBatch,
            adapter_id=adapter_id,
            adapter_source=adapter_source,
            revision=revision,
            **flash_causal_lm_kwargs,
        )

    if model_type == "opt":
        return OPTSharded(
            model_id,
            revision,
            quantize=quantize,
            compile=compile,
            dtype=dtype,
            trust_remote_code=trust_remote_code,
        )

    if model_type == "t5":
        return T5Sharded(
            model_id,
            revision,
            quantize=quantize,
            compile=compile,
            dtype=dtype,
            trust_remote_code=trust_remote_code,
        )

    if model_type == "granite":
        from lorax_server.models.flash_granite import FlashGranite

        return FlashGranite(
            model_id,
            adapter_id,
            adapter_source,
            revision,
            **flash_causal_lm_kwargs,
        )
    
    if model_type == "exaone":
        from lorax_server.models.flash_exaone import FlashExaOne

        return FlashExaOne(
            model_id,
            adapter_id,
            adapter_source,
            revision,
            **flash_causal_lm_kwargs,
        )

    raise ValueError(f"Unsupported model type {model_type}")