Support SDAR

[grimoire]

Since we are refactoring chat template, we will use template of internlm2 for now.

from lmdeploy import pipeline, PytorchEngineConfig, GenerationConfig, ChatTemplateConfig


from lmdeploy.pytorch.tools.utils import Timer, visualize_pipe_out


if __name__ == '__main__':
    model_path = 'JetLM/SDAR-1.7B-Chat'
    chat_template_config = ChatTemplateConfig('internlm2')

    log_level = 'WARNING'

    dllm_unmasking_strategy='low_confidence_dynamic'
    # dllm_unmasking_strategy='sequential'

    prompts = [
        'hakuna matata!',
        'The quick brown fox jumps over the lazy dog.'
        ]

    backend_config = PytorchEngineConfig(
        tp=1,
        dllm_block_length=4,
        dllm_unmasking_strategy=dllm_unmasking_strategy,
    )

    gen_config = GenerationConfig(
        max_new_tokens=512,
    )


    with pipeline(model_path, backend_config=backend_config,
                  chat_template_config=chat_template_config, log_level=log_level) as pipe:
        outputs = pipe(prompts, gen_config=gen_config)
        print(outputs)