VictorSanh's picture
Update visualization
217780a
raw
history blame
1.04 kB
from m4.models.custom_modules import DecoupledEmbedding, DecoupledLinear
from m4.models.vbloom.configuration_vbloom import VBloomConfig
from m4.models.vbloom.modeling_vbloom import VBloomForCausalLM
from m4.models.vgpt2.configuration_vgpt2 import VGPT2Config
from m4.models.vgpt2.modeling_vgpt2 import VGPT2LMHeadModel
from m4.models.vllama.configuration_vllama import VLlamaConfig
from m4.models.vllama.modeling_vllama import VLlamaForCausalLM
from m4.models.vopt.configuration_vopt import VOPTConfig
from m4.models.vopt.modeling_vopt import VOPTForCausalLM
from m4.models.vt5.configuration_vt5 import VT5Config
from m4.models.vt5.modeling_vt5 import VT5ForConditionalGeneration
_SUPPORTED_MODELS = {
"vgpt2": VGPT2Config,
"vt5": VT5Config,
"vbloom": VBloomConfig,
"vopt": VOPTConfig,
"vllama": VLlamaConfig,
}
model_type_to_modeling_class = {
"vgpt2": VGPT2LMHeadModel,
"vt5": VT5ForConditionalGeneration,
"vbloom": VBloomForCausalLM,
"vopt": VOPTForCausalLM,
"vllama": VLlamaForCausalLM,
}