Add API wrapper

app.py

@@ -12,8 +12,10 @@ fix_pytorch_int8()
 
 
 import torch
+import psutil
 import logging
 import gradio as gr
+from threading import Thread
 from transformers import AutoTokenizer, GenerationConfig, AutoModel
 
 
@@ -52,6 +54,28 @@ logging.basicConfig(
     level=logging.INFO,
     datefmt='%m/%d %H:%M:%S')
 
+
+def log_sys_info():
+    cpu_cores = psutil.cpu_count()
+    cpu_freq = '{:.2f}'.format(psutil.cpu_freq().max / 1000) + 'GHz'
+    mem = psutil.virtual_memory()
+    mem_total = '{:.2f}'.format(mem.total / 1024 / 1024 / 1024) + 'GB'
+    mem_used = '{:.2f}'.format(mem.used / 1024 / 1024 / 1024) + 'GB'
+    mem_percent = '{:.2f}'.format(mem.percent) + '%'
+    disk = psutil.disk_usage('.')
+    disk_total = '{:.2f}'.format(disk.total / 1024 / 1024 / 1024) + 'GB'
+    disk_used = '{:.2f}'.format(disk.used / 1024 / 1024 / 1024) + 'GB'
+    disk_percent = '{:.2f}'.format(disk.percent) + '%'
+
+    logging.info('======== SYSTEM INFO =========')
+    logging.info(f'CPU: {cpu_cores} cores, {cpu_freq}')
+    logging.info(f'RAM: {mem_used} / {mem_total}, {mem_percent} used')
+    logging.info(f'DISK: {disk_used} / {disk_total}, {disk_percent} used')
+    logging.info('==============================')
+
+
+log_sys_info()
+
 model = AutoModel.from_pretrained(
     int8_model,
     trust_remote_code=True,
@@ -67,6 +91,38 @@ model.eval()
 
 torch.set_default_tensor_type(torch.FloatTensor)
 
+logging.info('[SYS] Model loaded')
+log_sys_info()
+
+
+class CHAT_DB:
+    def __init__(self):
+        self.prompts = {}
+        self.results = {}
+        self.index = 1
+        self.lock = False
+    
+    def set(self, index, prompt=None, result=None):
+        assert prompt or result
+        if prompt:
+            if index in self.prompts:
+                raise ValueError('Prompt already exists')
+            self.prompts[index] = prompt
+            index += 1
+        if result:
+            self.results[index] = result
+    
+    def clean(self):
+        if len(self.prompts) > 100:
+            self.prompts = dict(list(self.prompts.items())[-100:])
+        k = list(set(self.prompts.keys()).intersection(set(self.results.keys())))  # keys to preserve
+        self.prompts = {i: self.prompts[i] for i in k}
+        self.results = {i: self.results[i] for i in k}
+        log_sys_info()
+
+
+db = CHAT_DB()
+
 
 def evaluate(context, temperature, top_p):
     generation_config = GenerationConfig(
@@ -99,6 +155,64 @@ def evaluate(context, temperature, top_p):
         return out_text
 
 
+def evaluate_wrapper(context, temperature, top_p):
+    db.lock = True
+    index = db.index
+    db.set(index, prompt=context)
+    result = evaluate(context, temperature, top_p)
+    db.set(index, result=result)
+    db.lock = False
+    return result
+
+
+def api_wrapper(context='', temperature=0.5, top_p=0.8, query=0):
+    query = int(query)
+    assert context or query
+    
+    return_json = {
+        'status': '',
+        'code': 0,
+        'message': '',
+        'index': 0,
+        'result': ''
+    }
+
+    if context:
+        if db.lock:
+            logging.info(f'[API] Request: {context}, Status: busy')
+            return_json['status'] = 'busy'
+            return_json['code'] = 503
+            return_json['message'] = 'Server is busy, please try again later.'
+            return return_json
+        else:
+            index = db.index
+            t = Thread(target=evaluate_wrapper, args=(context, temperature, top_p))
+            t.start()
+            logging.info(f'[API] Request: {context}, Status: processing, Index: {index}')
+            return_json['status'] = 'processing'
+            return_json['code'] = 202
+            return_json['message'] = 'Request accepted, please check back later.'
+            return_json['index'] = index
+            return return_json
+    else:  # query
+        if query in db.prompts:
+            if query in db.results:
+                logging.info(f'[API] Query: {query}, Status: hit')
+                return_json['status'] = 'done'
+                return_json['code'] = 200
+                return_json['message'] = 'Request processed.'
+                return_json['index'] = query
+                return_json['result'] = db.results[query]
+                return return_json
+            else:
+                logging.info(f'[API] Query: {query}, Status: processing')
+                return_json['status'] = 'processing'
+                return_json['code'] = 202
+                return_json['message'] = 'Request accepted, please check back later.'
+                return_json['index'] = query
+                return return_json
+
+
 def evaluate_stream(msg, history, temperature, top_p):
     generation_config = GenerationConfig(
         temperature=temperature,
@@ -158,12 +272,15 @@ with gr.Blocks() as demo:
             msg = gr.Textbox(label="输入框", placeholder="最近过得怎么样？",
                 info="输入你的内容，按 [Enter] 发送。什么都不填经常会出错。")
             clear = gr.Button("清除聊天")
-            api_handler = gr.Button("API", visible=False)
-            textbox_for_api = gr.Textbox(visible=False)
+
+        api_handler = gr.Button("API", visible=False)
+        num_for_api = gr.Number(visible=False)
+        json_for_api = gr.JSON(visible=False)
+
 
     msg.submit(evaluate_stream, [msg, chatbot, temp, top_p], [chatbot, msg])
     clear.click(lambda: None, None, chatbot, queue=False)
-    api_handler.click(evaluate, [textbox_for_api, temp, top_p], [textbox_for_api], api_name='chat')
+    api_handler.click(api_wrapper, [msg, temp, top_p, num_for_api], [json_for_api], api_name='chat')
     gr.HTML(gr_footer)
 
 demo.queue()

requirements.txt

@@ -1,3 +1,5 @@
+psutil
+
 # https://huggingface.co/spaces/ljsabc/Fujisaki/blob/main/requirements.txt
 
 # int8