feature(nyz): add naive Mistral-7B-Instruct-v0.1 demo

README.md

@@ -29,13 +29,9 @@ QUESTION_LANG=cn QUESTION_LLM='chatgpt' QUESTION_LLM_KEY=<your API key> python3
 ```shell
 QUESTION_LANG=en QUESTION_LLM='chatgpt' QUESTION_LLM_KEY=<your API key> python3 -u app.py
 ```
-### LLaMA2-7b + 中文
+### Mistral-7B-Instruct-v0.1 + 英文
 ```shell
-QUESTION_LANG=cn QUESTION_LLM='llama2-7b' python3 -u app.py
-```
-### LLaMA2-7b + 英文
-```shell
-QUESTION_LANG=en QUESTION_LLM='llama2-7b' python3 -u app.py
+QUESTION_LANG=en QUESTION_LLM='mistral-7b' python3 -u app.py
 ```
 ## :technologist: 为什么制作这个游戏
 
@@ -50,9 +46,9 @@ QUESTION_LANG=en QUESTION_LLM='llama2-7b' python3 -u app.py
 - [x] 支持自定义关卡
 - [ ] 在线试玩链接
 - [ ] Hugging Face Space 链接
-- [ ] 支持LLaMA2-7B（英文）
-- [ ] 支持Mistral-7B（英文）
+- [x] 支持Mistral-7B-Instruct-v0.1（英文）
 - [ ] 支持Baichuan2-7B（中文）
+- [ ] 支持LLaMA2-7B（英文）
 - [ ] LLM 推理速度优化
 
 

app.py

@@ -11,7 +11,7 @@ _QUESTIONS = list_ordered_questions()
 _LANG = os.environ.get('QUESTION_LANG', 'cn')
 assert _LANG in ['cn', 'en'], _LANG
 _LLM = os.environ.get('QUESTION_LLM', 'chatgpt')
-assert _LLM in ['chatgpt', 'llama2-7b'], _LLM
+assert _LLM in ['chatgpt', 'mistral-7b'], _LLM
 _LLM_KEY = os.environ.get('QUESTION_LLM_KEY', None)
 
 if _LANG == "cn":

llmriddles/llms/__init__.py

@@ -1,2 +1,3 @@
-from .chatgpt import ask_chatgpt
 from .base import register_llm, get_llm_fn
+from .chatgpt import ask_chatgpt
+from .mistral import ask_mistral_7b_instruct

llmriddles/llms/llm_client.py

@@ -0,0 +1,54 @@
+import time
+import requests
+import logging
+import argparse
+
+
+class LLMFlaskClient:
+    def __init__(self, ip: str, port: int, max_retry: int = 3):
+        self.ip = ip
+        self.port = port
+
+        self.url_prefix_format = 'http://{}:{}/'
+        self.url = self.url_prefix_format.format(self.ip, self.port)
+        self.max_retry = max_retry
+
+        self.logger = logging.getLogger()
+        self.logger.addHandler(logging.StreamHandler())
+        self.logger.handlers[0].setFormatter(logging.Formatter("%(message)s"))
+
+    def _request(self, name: str, data: dict):
+        for _ in range(self.max_retry):
+            try:
+                self.logger.info(f'{name}\ndata: {data}')
+                response = requests.post(self.url + name, json=data).json()
+            except Exception as e:
+                self.logger.warning('error: ', repr(e))
+                time.sleep(1)
+                continue
+            if response['code'] == 0:
+                return response['output']
+            else:
+                raise Exception(response['error_msg'])
+        raise Exception("Web service failed. Please retry or contact with manager")
+
+    def run(self, message: str) -> str:
+        try:
+            return self._request('ask_llm_for_answer', {'user_text': message})
+        except Exception as e:
+            return f"Error: {repr(e)}"
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--ip', required=True)
+    parser.add_argument('-p', '--port', required=True)
+    parser.add_argument('--debug', action='store_true')
+    args = parser.parse_args()
+    if args.debug:
+        logging.getLogger().setLevel(logging.INFO)
+    else:
+        logging.getLogger().setLevel(logging.WARNING)
+
+    client = LLMFlaskClient(args.ip, args.port)
+    print(client.run('Please concatenate string "1+" and "1=3". Only give me the result without "".'))

llmriddles/llms/llm_server.py

@@ -0,0 +1,72 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from flask import Flask, request
+import argparse
+import logging
+
+
+class LLMInstance:
+
+    def __init__(self, model_path: str, device: str = "cuda"):
+
+        self.model = AutoModelForCausalLM.from_pretrained(model_path)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model.to(device)
+        self.device = device
+
+    def query(self, message):
+        try:
+            messages = [
+                {"role": "user", "content": message},
+            ]
+            encodeds = self.tokenizer.apply_chat_template(messages, return_tensors="pt")
+            model_inputs = encodeds.to(self.device)
+
+            generated_ids = self.model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
+            decoded = self.tokenizer.batch_decode(generated_ids)
+
+            # output is the string decoded[0] after "[/INST]". There may exist "</s>", delete it.
+            output = decoded[0].split("[/INST]")[1].split("</s>")[0]
+            return {
+                'code': 0,
+                'ret': True,
+                'error_msg': None,
+                'output': output
+            }
+        except Exception as e:
+            return {
+                'code': 1,
+                'ret': False,
+                'error_msg': str(e),
+                'output': None
+            }
+
+
+def create_app(core):
+    app = Flask(__name__)
+
+    @app.route('/ask_llm_for_answer', methods=['POST'])
+    def ask_llm_for_answer():
+        user_text = request.json['user_text']
+        return core.query(user_text)
+
+    return app
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-m', '--model_path', required=True, default='Mistral-7B-Instruct-v0.1', help='the model path of reward model')
+    parser.add_argument('--ip', default='0.0.0.0')
+    parser.add_argument('-p', '--port', default=8001)
+    parser.add_argument('--debug', action='store_true')
+    args = parser.parse_args()
+
+    if args.debug:
+        logging.getLogger().setLevel(logging.DEBUG)
+    else:
+        logging.getLogger().setLevel(logging.INFO)
+    logging.getLogger().addHandler(logging.StreamHandler())
+    logging.getLogger().handlers[0].setFormatter(logging.Formatter("%(message)s"))
+
+    core = LLMInstance(args.model_path)
+    app = create_app(core)
+    app.run(host=args.ip, port=args.port)

llmriddles/llms/mistral.py

@@ -0,0 +1,22 @@
+from functools import lru_cache
+
+from .base import register_llm
+from .llm_client import LLMFlaskClient
+
+
+@lru_cache()
+def _get_mistral_7b_instruct_server(host: str, port: int):
+    from .llm_server import LLMInstance, create_app
+    core = LLMInstance('Mistral-7B-Instruct-v0.1')
+    app = create_app(core)
+    app.run(host=host, port=port)
+
+
+def ask_mistral_7b_instruct(message: str, **kwargs):
+    host, port = '0.0.0.0', 8001
+    _get_mistral_7b_instruct_server(host, port)
+    client = LLMFlaskClient(host, port)
+    return client.run(message).strip()
+
+
+register_llm('mistral-7b', ask_mistral_7b_instruct)

requirements.txt

@@ -2,4 +2,6 @@ hbutils>=0.9.1
 tqdm
 requests>=2.20
 gradio==4.1.1
-openai>=1
+openai>=1
+flask
+transformers