Class LocalModel

src/model.py:159–309

class LocalModel(object):
    def __init__(self, model, device_num, llama = False, tensor_parallel_size = 4, temperature=0.7, n = 200) -> None:
        self.name = model
        self.sampling_params = SamplingParams(temperature=temperature, n = n, max_tokens = 512)
        #self.sampling_params = SamplingParams(temperature=0, n = 1, max_tokens = 512)
        self.model = LLM(model, trust_remote_code = True, tensor_parallel_size= tensor_parallel_size)
        self.tokenizer = self.model.get_tokenizer()

    def get_prompt_length(self, prompt):
        return len(self.tokenizer(prompt).input_ids)


    def format_chats(self, dialogs, pad = False):
        B_INST, E_INST = "[INST]", "[/INST]"
        B_SYS, E_SYS = "SYS\n", "\n<</SYS>>\n\n"


        prompt_tokens = []


        for dialog in dialogs:
            if dialog[0]["role"] == "system":
                dialog = [
                    {
                        "role": dialog[1]["role"],
                        "content": B_SYS
                        + dialog[0]["content"]
                        + E_SYS
                        + dialog[1]["content"],
                    }
                ] + dialog[2:]
            
            assert all([msg["role"] == "user" for msg in dialog[::2]]) and all(
                    [msg["role"] == "assistant" for msg in dialog[1::2]]
                ), (
                    "model only supports 'system', 'user' and 'assistant' roles, "
                    "starting with 'system', then 'user' and alternating (u/a/u/a/u...)"
                )

            dialog_tokens = sum(
                [
                    [
                        [self.tokenizer.bos_token_id]
                        + self.tokenizer(
                            f"{B_INST} {(prompt['content']).strip()} {E_INST} {(answer['content']).strip()} "
                        ).input_ids
                        + [self.tokenizer.eos_token_id]
                    ]
                    for prompt, answer in zip(
                        dialog[::2],
                        dialog[1::2],
                    )
                ],
                [],
            )
            assert (
                dialog[-1]["role"] == "user"
            ), f"Last message must be from user, got {dialog[-1]['role']}"