You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
16 lines
580 B
Python
16 lines
580 B
Python
from petals.models.mixtral.block import WrappedMixtralBlock
|
|
from petals.models.mixtral.config import DistributedMixtralConfig
|
|
from petals.models.mixtral.model import (
|
|
DistributedMixtralForCausalLM,
|
|
DistributedMixtralForSequenceClassification,
|
|
DistributedMixtralModel,
|
|
)
|
|
from petals.utils.auto_config import register_model_classes
|
|
|
|
register_model_classes(
|
|
config=DistributedMixtralConfig,
|
|
model=DistributedMixtralModel,
|
|
model_for_causal_lm=DistributedMixtralForCausalLM,
|
|
model_for_sequence_classification=DistributedMixtralForSequenceClassification,
|
|
)
|