mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
12 lines
360 B
12 lines
360 B
from .bloom import BloomModelInferPolicy
|
|
from .chatglm2 import ChatGLM2InferPolicy
|
|
from .llama import LlamaModelInferPolicy
|
|
|
|
model_policy_map = {
|
|
"llama": LlamaModelInferPolicy,
|
|
"bloom": BloomModelInferPolicy,
|
|
"chatglm": ChatGLM2InferPolicy,
|
|
}
|
|
|
|
__all__ = ["LlamaModelInferPolicy", "BloomModelInferPolicy", "ChatGLM2InferPolicy", "model_polic_map"]
|