add chatbot (InternLM#2)

.pre-commit-config.yaml

@@ -50,4 +50,4 @@ repos:
     rev: v0.2.0
     hooks:
     -   id: check-copyright
-        args: []
+        args: ["llmdeploy"]

llmdeploy/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) OpenMMLab. All rights reserved.

llmdeploy/serve/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) OpenMMLab. All rights reserved.

llmdeploy/serve/client.py

@@ -0,0 +1,38 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import os
+
+import fire
+
+from llmdeploy.serve.fastertransformer.chatbot import Chatbot
+
+
+def input_prompt():
+    print('\ndouble enter to end input >>> ', end='')
+    sentinel = ''  # ends when this string is seen
+    return '\n'.join(iter(input, sentinel))
+
+
+def main(triton_server_addr: str, model_name: str, session_id: int):
+    log_level = os.environ.get('SERVICE_LOG_LEVEL', 'INFO')
+    chatbot = Chatbot(triton_server_addr,
+                      model_name,
+                      log_level=log_level,
+                      display=True)
+    nth_round = 1
+    while True:
+        prompt = input_prompt()
+        if prompt == 'exit':
+            exit(0)
+        elif prompt == 'end':
+            chatbot.end(session_id)
+        else:
+            request_id = f'{session_id}-{nth_round}'
+            for status, res, tokens in chatbot.stream_infer(
+                    session_id, prompt, request_id=request_id):
+                continue
+            print(f'session {session_id}, {status}, {tokens}, {res}')
+        nth_round += 1
+
+
+if __name__ == '__main__':
+    fire.Fire(main)

llmdeploy/serve/fastertransformer/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from llmdeploy.serve.fastertransformer.chatbot import \
+    Chatbot  # noqa: F401,F403