

def load_llm(args):
    if "openai" in args.model:
        from models.openai import ChatGPT
        llm = ChatGPT(model_path=args.model, max_tokens=args.max_token)
    else:
        from models.hf import CasualLM
        llm = CasualLM(model_path=args.model, arch=args.arch,
                       max_tokens=args.max_token, infer_mode=args.infer_mode)
    return llm

