Feat: support stream output in cli_demo (#137)

* Feat: support stream output in cli_demo

* support generation in cli_demo

examples/model_cli_demo.py

@@ -11,6 +11,11 @@ def parse_args():
         type=str,
         default='internlm/internlm2-chat-20b',
         help='The path to the model')
+    parser.add_argument(
+        '--mode',
+        type=str,
+        default='chat',
+        help='Completion through chat or generate')
     args = parser.parse_args()
     return args
 
@@ -42,9 +47,16 @@ def main():
         if prompt == 'exit':
             exit(0)
         history.append(dict(role='user', content=prompt))
-        response = model.chat(history, max_new_tokens=512)
+        if args.mode == 'generate':
+            history = [dict(role='user', content=prompt)]
+        print('\nInternLm2：', end='')
+        current_length = 0
+        for status, response, _ in model.stream_chat(
+                history, max_new_tokens=512):
+            print(response[current_length:], end='', flush=True)
+            current_length = len(response)
         history.append(dict(role='assistant', content=response))
-        print('Assistant:', response)
+        print('')
 
 
 if __name__ == '__main__':