mirror of https://github.com/hpcaitech/ColossalAI
12 lines
360 B
Python
12 lines
360 B
Python
|
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"]
|