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
556 B
Python
16 lines
556 B
Python
from petals.models.bloom.block import WrappedBloomBlock
|
|
from petals.models.bloom.config import DistributedBloomConfig
|
|
from petals.models.bloom.model import (
|
|
DistributedBloomForCausalLM,
|
|
DistributedBloomForSequenceClassification,
|
|
DistributedBloomModel,
|
|
)
|
|
from petals.utils.auto_config import register_model_classes
|
|
|
|
register_model_classes(
|
|
config=DistributedBloomConfig,
|
|
model=DistributedBloomModel,
|
|
model_for_causal_lm=DistributedBloomForCausalLM,
|
|
model_for_sequence_classification=DistributedBloomForSequenceClassification,
|
|
)
|