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}'