modulize

app.py

@@ -2,28 +2,18 @@ from fix_int8 import fix_pytorch_int8
 fix_pytorch_int8()
 
 
-# import subprocess
-# result = subprocess.run(['git', 'clone', 'https://huggingface.co/KumaTea/twitter-int8', 'model'], capture_output=True, text=True)
-# print(result.stdout)
-
-
 # Credit:
 # https://huggingface.co/spaces/ljsabc/Fujisaki/blob/main/app.py
 
 
 import torch
-import psutil
-import logging
 import gradio as gr
 from threading import Thread
+from model import model, tokenizer
+from session import db, logger, log_sys_info
 from transformers import AutoTokenizer, GenerationConfig, AutoModel
 
 
-chatglm = 'THUDM/chatglm-6b'
-chatglm_rev = '4de8efe'
-int8_model = 'KumaTea/twitter-int8'
-int8_model_rev = '1136001'
-
 max_length = 224
 default_start = ["你是Kuma，请和我聊天，每句话以两个竖杠分隔。", "好的，你想聊什么？"]
 
@@ -45,85 +35,6 @@ gr_footer =  """<p align='center'>
 </p>"""
 
 
-
-# device = torch.device('cpu')
-# torch.cuda.current_device = lambda : device
-
-logging.basicConfig(
-    format='%(asctime)s %(levelname)-8s %(message)s',
-    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,
-    revision=int8_model_rev
-).float()  # .to(device)
-tokenizer = AutoTokenizer.from_pretrained(chatglm, trust_remote_code=True, revision=chatglm_rev)
-
-# dump a log to ensure everything works well
-# print(model.peft_config)
-# We have to use full precision, as some tokens are >65535
-model.eval()
-# print(model)
-
-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(
         temperature=temperature,
@@ -139,7 +50,7 @@ def evaluate(context, temperature, top_p):
         # No need for starting prompt in API
         if not context.endswith('||'):
             context += '||'
-        logging.info('[API] Request: ' + context)
+        # logger.info('[API] Request: ' + context)
         ids = tokenizer([context], return_tensors="pt")
         inputs = ids.to("cpu")
         out = model.generate(
@@ -151,17 +62,17 @@ def evaluate(context, temperature, top_p):
         decoder_output = tokenizer.decode(out)
         # out_text = decoder_output.split("Answer: ")[1]
         out_text = decoder_output
-        logging.info('[API] Results: ' + out_text)
+        logger.info('[API] Results: ' + out_text.replace('\n', '<br>'))
         return out_text
 
 
 def evaluate_wrapper(context, temperature, top_p):
-    db.lock = True
+    db.lock()
     index = db.index
     db.set(index, prompt=context)
     result = evaluate(context, temperature, top_p)
     db.set(index, result=result)
-    db.lock = False
+    db.unlock()
     return result
 
 
@@ -178,37 +89,53 @@ def api_wrapper(context='', temperature=0.5, top_p=0.8, query=0):
     }
 
     if context:
-        if db.lock:
-            logging.info(f'[API] Request: {context}, Status: busy')
+        if db.islocked():
+            logger.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_json['message'] = '[context] Server is busy, please try again later.'
             return return_json
         else:
+            for index in db.prompts:
+                if db.prompts[index] == context:
+                    return_json['status'] = 'done'
+                    return_json['code'] = 200
+                    return_json['message'] = '[context] Request cached.'
+                    return_json['index'] = index
+                    return_json['result'] = db.results[index]
+                    return return_json
+            # new
             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}')
+            logger.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['message'] = '[context] 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.'
+        if query in db.prompts and query in db.results:
+            logger.info(f'[API] Query: {query}, Status: hit')
+            return_json['status'] = 'done'
+            return_json['code'] = 200
+            return_json['message'] = '[query] Request processed.'
+            return_json['index'] = query
+            return_json['result'] = db.results[query]
+            return return_json
+        else:
+            if db.islocked():
+                logger.info(f'[API] Query: {query}, Status: processing')
+                return_json['status'] = 'processing'
+                return_json['code'] = 202
+                return_json['message'] = '[query] Request in processing, please check back later.'
                 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.'
+                logger.info(f'[API] Query: {query}, Status: error')
+                return_json['status'] = 'error'
+                return_json['code'] = 404
+                return_json['message'] = '[query] Index not found.'
                 return_json['index'] = query
                 return return_json
 
@@ -247,11 +174,11 @@ def evaluate_stream(msg, history, temperature, top_p):
         context = context[15:]
 
     h = []
-    logging.info('[UI] Request: ' + context)
+    logger.info('[UI] Request: ' + context)
     for response, h in model.stream_chat(tokenizer, context, h, max_length=max_length, top_p=top_p, temperature=temperature):
         history[-1][1] = response
         yield history, ""
-    logging.info('[UI] Results: ' + response)
+    logger.info('[UI] Results: ' + response.replace('\n', '<br>'))
 
 
 with gr.Blocks() as demo:
@@ -274,13 +201,16 @@ with gr.Blocks() as demo:
             clear = gr.Button("清除聊天")
 
         api_handler = gr.Button("API", visible=False)
-        num_for_api = gr.Number(visible=False)
-        json_for_api = gr.JSON(visible=False)
+        api_index = gr.Number(visible=False)
+        api_result = gr.JSON(visible=False)
+        info_handler = gr.Button("Info", visible=False)
+        info_text = gr.Textbox('System info logged. Check it in the log viewer.', visible=False)
 
 
     msg.submit(evaluate_stream, [msg, chatbot, temp, top_p], [chatbot, msg])
     clear.click(lambda: None, None, chatbot, queue=False)
-    api_handler.click(api_wrapper, [msg, temp, top_p, num_for_api], [json_for_api], api_name='chat')
+    api_handler.click(api_wrapper, [msg, temp, top_p, api_index], api_result, api_name='chat')
+    info_handler.click(log_sys_info, None, info_text, api_name='info')
     gr.HTML(gr_footer)
 
 demo.queue()

fix_int8.py

@@ -1,5 +1,6 @@
 import os
 import sys
+from session import logger
 
 
 def fix_pytorch_int8():
@@ -26,4 +27,4 @@ def fix_pytorch_int8():
         with open(fix_path, 'w') as f:
             f.write(text)
 
-        return print('Fixed torch/nn/parameter.py')
+        return logger.info('Fixed torch/nn/parameter.py')

model.py

@@ -0,0 +1,36 @@
+import torch
+from session import logger, log_sys_info
+from transformers import AutoTokenizer, GenerationConfig, AutoModel
+
+
+chatglm = 'THUDM/chatglm-6b'
+chatglm_rev = '4de8efe'
+int8_model = 'KumaTea/twitter-int8'
+int8_model_rev = '1136001'
+
+# import subprocess
+# result = subprocess.run(['git', 'clone', 'https://huggingface.co/KumaTea/twitter-int8', 'model'], capture_output=True, text=True)
+# print(result.stdout)
+
+# device = torch.device('cpu')
+# torch.cuda.current_device = lambda : device
+
+log_sys_info()
+
+model = AutoModel.from_pretrained(
+    int8_model,
+    trust_remote_code=True,
+    revision=int8_model_rev
+).float()  # .to(device)
+tokenizer = AutoTokenizer.from_pretrained(chatglm, trust_remote_code=True, revision=chatglm_rev)
+
+# dump a log to ensure everything works well
+# print(model.peft_config)
+# We have to use full precision, as some tokens are >65535
+model.eval()
+# print(model)
+
+torch.set_default_tensor_type(torch.FloatTensor)
+
+logger.info('[SYS] Model loaded')
+log_sys_info()

session.py

@@ -0,0 +1,71 @@
+import os
+import psutil
+import logging
+from pathlib import Path
+
+
+logging.basicConfig(
+    format='%(asctime)s %(levelname)-8s %(message)s',
+    level=logging.INFO,
+    datefmt='%m/%d %H:%M:%S')
+logger = logging.getLogger(__name__)
+
+
+def log_sys_info():
+    cpu_cores = psutil.cpu_count()
+    # cpu_freq = '{:.2f}'.format(psutil.cpu_freq().max / 1000) + 'GHz'
+    cpu_percent = '{:.2f}'.format(psutil.cpu_percent()) + '%'
+    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) + '%'
+
+    logger.info('======== SYSTEM INFO =========')
+    logger.info(f'CPU: {cpu_cores} cores, {cpu_percent} used')
+    logger.info(f'RAM: {mem_used} / {mem_total}, {mem_percent} used')
+    logger.info(f'DISK: {disk_used} / {disk_total}, {disk_percent} used')
+    logger.info('==============================')
+
+
+class CHAT_DB:
+    def __init__(self):
+        self.prompts = {}
+        self.results = {}
+        self.index = 1
+        self.lockfile = '.lock'
+
+    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
+            self.index += 1
+        if result:
+            self.results[index] = result
+
+    def lock(self):
+        if not os.path.exists(self.lockfile):
+            Path(self.lockfile).touch(exist_ok=True)
+
+    def unlock(self):
+        if os.path.exists(self.lockfile):
+            os.remove(self.lockfile)
+
+    def islocked(self):
+        return os.path.exists(self.lockfile)
+
+    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()