CCCasEEE/internlm_lagent

1from argparse import ArgumentParser2 3from lagent.llms import HFTransformer4from lagent.llms.meta_template import INTERNLM2_META as META5 6 7def parse_args():8    parser = ArgumentParser(description='chatbot')9    parser.add_argument(10        '--path',11        type=str,12        default='internlm/internlm2-chat-20b',13        help='The path to the model')14    parser.add_argument(15        '--mode',16        type=str,17        default='chat',18        help='Completion through chat or generate')19    args = parser.parse_args()20    return args21 22 23def main():24    args = parse_args()25    # Initialize the HFTransformer-based Language Model (llm)26    model = HFTransformer(27        path=args.path,28        meta_template=META,29        max_new_tokens=1024,30        top_p=0.8,31        top_k=None,32        temperature=0.1,33        repetition_penalty=1.0,34        stop_words=['<|im_end|>'])35 36    def input_prompt():37        print('\ndouble enter to end input >>> ', end='', flush=True)38        sentinel = ''  # ends when this string is seen39        return '\n'.join(iter(input, sentinel))40 41    history = []42    while True:43        try:44            prompt = input_prompt()45        except UnicodeDecodeError:46            print('UnicodeDecodeError')47            continue48        if prompt == 'exit':49            exit(0)50        history.append(dict(role='user', content=prompt))51        if args.mode == 'generate':52            history = [dict(role='user', content=prompt)]53        print('\nInternLm2：', end='')54        current_length = 055        for status, response, _ in model.stream_chat(history):56            print(response[current_length:], end='', flush=True)57            current_length = len(response)58        history.append(dict(role='assistant', content=response))59        print('')60 61 62if __name__ == '__main__':63    main()64