Class GetAssistantAns

src/getAssistantAns.py:7–41

from transformers import AutoModelForCausalLM

class GetAssistantAns():
    # 按照自己推理需求自己修改代码

    def __init__(self, gpu_num=1):
        model = AutoModelForCausalLM.from_pretrained(model_name)
        device_list = []
        for gpu_idx in range(gpu_num):
            device_list.append(torch.device("cuda:0"))

        # 将模型移动到指定的GPU设备
        model.to(device)


    def gen_answer(self, chat_dict, gpu_index):
        # 这里实际根据自己推理逻辑 然后转为标准格式返回
        # 以下仅仅是样例
        import time
        print(os.environ["CUDA_VISIBLE_DEVICES"])
        time.sleep(1)
        rtn_dict1 = {
                "role": "assistant",
                "content": None,
                "function_call":
                {
                    "name": "get_fudan_university_scoreline",
                    "arguments": "{\n  \"year\": \"2020\"\n}"
                }
            }

        rtn_dict2 =  {
                "role": "assistant",
                "content": "2020年复旦大学的分数线如下：\n\n- 文科一批：630分\n- 文科二批：610分\n- 理科一批：650分\n- 理科二批：630分"
            }

        return random.choice([rtn_dict1, rtn_dict2])

# ======================================================================
# 下面注释的部分是一个huggingface推理的多卡的demo