| 74 | return input_ids |
| 75 | |
| 76 | class GetAssistantAns(): |
| 77 | # 按照自己推理需求自己修改代码 |
| 78 | |
| 79 | def __init__(self): |
| 80 | pass |
| 81 | |
| 82 | def gen_answer(self, chat_dict): |
| 83 | chatrounds_list = change2chatml(chat_dict) |
| 84 | input_ids = get_chatrounds_ids(chatrounds_list) |
| 85 | output_ids = model.generate(torch.tensor([input_ids]).to(model.device), max_new_tokens=output_len, num_beams=1, num_return_sequences=1, do_sample=True, temperature=temperature, top_p=top_p, eos_token_id=end_token_id, top_k=top_k, streamer=None, repetition_penalty=repetition_penalty, pad_token_id=10000)[0] |
| 86 | res = tokenizer.decode(output_ids[len(input_ids):-1]) |
| 87 | save_dict = {"role": "assistant"} |
| 88 | if res.startswith("#function"): |
| 89 | try: |
| 90 | res_dict = json.loads(re.sub("^#function", "", res)) |
| 91 | save_dict["content"] = res_dict["content"] |
| 92 | save_dict["function_call"] = {} |
| 93 | save_dict["function_call"]["name"] = res_dict["name"] |
| 94 | save_dict["function_call"]["arguments"] = res_dict["arguments"] |
| 95 | except Exception as e: |
| 96 | print(e) |
| 97 | save_dict = {"role": "assistant"} |
| 98 | save_dict["content"] = res |
| 99 | else: |
| 100 | save_dict["content"] = res |
| 101 | |
| 102 | print(save_dict) |
| 103 | |
| 104 | return save_dict |
nothing calls this directly
no outgoing calls
no test coverage detected