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
568 B
Python
16 lines
568 B
Python
from petals.models.falcon.block import WrappedFalconBlock
|
|
from petals.models.falcon.config import DistributedFalconConfig
|
|
from petals.models.falcon.model import (
|
|
DistributedFalconForCausalLM,
|
|
DistributedFalconForSequenceClassification,
|
|
DistributedFalconModel,
|
|
)
|
|
from petals.utils.auto_config import register_model_classes
|
|
|
|
register_model_classes(
|
|
config=DistributedFalconConfig,
|
|
model=DistributedFalconModel,
|
|
model_for_causal_lm=DistributedFalconForCausalLM,
|
|
model_for_sequence_classification=DistributedFalconForSequenceClassification,
|
|
)
|