Update app.py

app.py

@@ -98,6 +98,9 @@ gpu_memory_utilization = float(os.environ.get("gpu_memory_utilization", "0.9"))
 # whether to enable quantization, currently not in use
 QUANTIZATION = str(os.environ.get("QUANTIZATION", ""))
 
+DATA_SET_REPO_PATH = str(os.environ.get("DATA_SET_REPO_PATH", ""))
+DATA_SET_REPO = None
+
 
 """
 Internal instructions of how to configure the DEMO
@@ -1207,6 +1210,34 @@ def format_conversation(history):
     return _str
 
 
+def maybe_upload_to_dataset():
+    global LOG_FILE, DATA_SET_REPO_PATH, SAVE_LOGS
+    if SAVE_LOGS and os.path.exists(LOG_PATH) and DATA_SET_REPO_PATH is not "":
+        with open(LOG_PATH, 'r', encoding='utf-8') as f:
+            convos = {}
+            for l in f:
+                if l:
+                    item = json.loads(l)
+                    convos[item['key']] = item
+        AGG_LOG_PATH = LOG_PATH + ".agg.json"
+        with open(AGG_LOG_PATH, 'w', encoding='utf-8') as fo:
+            json.dump(convos, fo, indent=4, ensure_ascii=False)
+        print(f'Saved aggregated json to {AGG_LOG_PATH}')
+        try:
+            from huggingface_hub import upload_file
+            print(f'upload {AGG_LOG_PATH} to {DATA_SET_REPO_PATH}')
+            upload_file(
+                path_or_fileobj=AGG_LOG_PATH,
+                path_in_repo=os.path.basename(AGG_LOG_PATH),
+                repo_id=DATA_SET_REPO_PATH,
+                token=HF_TOKEN,
+                repo_type="dataset",
+                create_pr=True
+            )
+        except Exception as e:
+            print(f'Failed to save to repo: {DATA_SET_REPO_PATH}|{str(e)}')
+            
+
 def print_log_file():
     global LOG_FILE, LOG_PATH
     if SAVE_LOGS and os.path.exists(LOG_PATH):
@@ -1222,6 +1253,7 @@ def print_log_file():
                 print(f'######--{v}--##')
                 _str = format_conversation(history)
                 print(_str)
+        maybe_upload_to_dataset()
 
 
 def debug_chat_response_echo(
@@ -1319,6 +1351,7 @@ def launch():
         f'\n| gpu_memory_utilization={gpu_memory_utilization} '
         f'\n| KEYWORDS={KEYWORDS} '
         f'\n| LOG_PATH={LOG_PATH} | SAVE_LOGS={SAVE_LOGS} '
+        f'\n| DATA_SET_REPO_PATH={DATA_SET_REPO_PATH} '
         f'\n| GET_LOG_CMD={GET_LOG_CMD} '
         f'\n| Sys={SYSTEM_PROMPT_1}'
         f'\n| Desc={model_desc}'