add self model to lm-evaluation-harness

Daze_Lu發表於2024-04-25

backgroud

I use huggingface transformers build a new moe model, when I use AutoForCasualModel to load the model, there is no suitable model structure to load it, in this case, the parameter couldn't be load correctly. To evaluate the performance of this model, I have to add a new style model into the lm-evaluation-harness.

build self model

https://github.com/learner-crapy/MOE-n-experts

add new model to lm-evaluation harness

click to check to code
# PhiMoe.py
from typing import Optional, Union

import torch

import lm_eval.models.utils
from lm_eval.api.registry import register_model
from lm_eval.models.huggingface import HFLM


@register_model("phimoe")
class PhiMoe(HFLM):
    def __init__(
        self,
        pretrained="/home/ludaze/Docker/Llama/MOE-n-experts/models/PhiMoeForCausalLM-0-1-top1",
        **kwargs,
    ) -> None:
        if "backend" in kwargs:
            # mamba currently only supports causal models
            assert kwargs["backend"] == "causal"

        super().__init__(
            pretrained=pretrained,
            # set appropriate defaults for tokenizer, max length, etc
            # backend=kwargs.pop("backend", "causal"),
            # tokenizer=kwargs.pop("tokenizer", "EleutherAI/gpt-neox-20b"),
            # max_length=kwargs.pop("max_length", 2048),
            **kwargs,
        )

    def _get_config(
        self,
        pretrained: str,
        **kwargs,
    ) -> None:
        try:
            from moe.transformers.models.phi.configuration_phi_moe import PhiMoeConfig
        except ModuleNotFoundError:
            raise Exception(
                "attempted to use 'PhiMoeForCausalLM' LM type, but package `phimoe` is not installed. \
please install phimoe via `pip install lm-eval[phimoe]` or `pip install -e .[phimoe]`",
            )

        self._config = PhiMoeConfig(pretrained)

    def _create_model(
        self,
        pretrained: str,
        # dtype: Optional[Union[str, torch.dtype]] = "float16",
        # no `parallelize=True` options
        # no PEFT and quantization options
        # Mamba does not support arbitrary HF from_pretrained() args
        **kwargs,
    ) -> None:
        try:
            from moe.transformers.models.phi.modeling_phi_moe import PhiMoeForCausalLM
        except ModuleNotFoundError:
            raise Exception(
                "attempted to use 'mamba_ssm' LM type, but package `phimoe` is not installed. \
please install phimoe via `pip install lm-eval[phimoe]` or `pip install -e .[phimoe]`",
            )

        self._model = PhiMoeForCausalLM.from_pretrained(pretrained)

    # def _model_generate(self, context, max_length, stop, **generation_kwargs):
    #     for key in ("do_sample", "attention_mask"):
    #         if key in generation_kwargs:
    #             generation_kwargs.pop(key)

        # return self.model.generate(
        #     input_ids=context,
        #     max_length=max_length,
        #     **generation_kwargs,
        # )

then register the file in lm_eval/models/init.py, PhiMoe.

lastly, run pip install -e .[phimoe]

相關文章