fix hf logging

model/cfg.yaml → cfg.yaml

@@ -1,4 +1,5 @@
-use_ckpt: hf
+launch_mode: hf   # hf or local
+hf_log_path: embodied-generalist/leo_demo_log
 hf_ckpt_path: [huangjy-pku/embodied-generalist, weights/leo_noact_hf.pth]
 local_ckpt_path: /mnt/huangjiangyong/leo/hf_assets/weights/leo_noact_lora.pth
 model:
@@ -6,7 +7,6 @@ model:
   # vision modules omitted
   llm:
     name: Vicuna7B
-    use_ckpt: hf
     hf_cfg_path: huangjy-pku/vicuna-7b
     local_cfg_path: /mnt/huangjiangyong/vicuna-7b
     truncation_side: right

model/leo_agent.py

@@ -14,16 +14,13 @@ def disabled_train(self, mode=True):
 class LeoAgentLLM(nn.Module):
     def __init__(self, cfg):
         super().__init__()
-        if hasattr(cfg, 'model'):
-            cfg = cfg.model
-
         # LLM
-        if cfg.llm.use_ckpt == 'hf':
-            llm_cfg_path = snapshot_download(cfg.llm.hf_cfg_path)
+        if cfg.launch_mode == 'hf':
+            llm_cfg_path = snapshot_download(cfg.model.llm.hf_cfg_path)
         else:
-            llm_cfg_path = cfg.llm.local_cfg_path
+            llm_cfg_path = cfg.model.llm.local_cfg_path
         self.llm_tokenizer = LlamaTokenizer.from_pretrained(llm_cfg_path, use_fast=False,
-                                                            truncation_side=cfg.llm.truncation_side)
+                                                            truncation_side=cfg.model.llm.truncation_side)
         self.llm_tokenizer.add_special_tokens({'pad_token': '[PAD]'})
         self.llm_tokenizer.add_special_tokens({'bos_token': '<s>'})
         self.llm_tokenizer.add_special_tokens({'eos_token': '</s>'})
@@ -37,18 +34,18 @@ class LeoAgentLLM(nn.Module):
         self.llm_model.train = disabled_train
 
         # LoRA-based LLM fine-tuning
-        if cfg.llm.lora.flag:
+        if cfg.model.llm.lora.flag:
             lora_config = LoraConfig(
-                r=cfg.llm.lora.rank,
-                lora_alpha=cfg.llm.lora.alpha,
-                target_modules=cfg.llm.lora.target_modules,
-                lora_dropout=cfg.llm.lora.dropout,
+                r=cfg.model.llm.lora.rank,
+                lora_alpha=cfg.model.llm.lora.alpha,
+                target_modules=cfg.model.llm.lora.target_modules,
+                lora_dropout=cfg.model.llm.lora.dropout,
                 bias='none',
                 modules_to_save=[],
             )
             self.llm_model = get_peft_model(self.llm_model, peft_config=lora_config)
 
-        self.max_context_len = cfg.llm.max_context_len
+        self.max_context_len = cfg.model.llm.max_context_len
 
     @property
     def device(self):

utils.py

@@ -3,11 +3,12 @@ import datetime
 import json
 import os
 import time
+from uuid import uuid4
 
 import gradio as gr
 import torch
 import yaml
-from huggingface_hub import hf_hub_download
+from huggingface_hub import CommitScheduler, hf_hub_download
 from omegaconf import OmegaConf
 
 from model.leo_agent import LeoAgentLLM
@@ -27,33 +28,34 @@ OBJECTS_PROMPT = "Objects (including you) in the scene:"
 TASK_PROMPT = "USER: {instruction} ASSISTANT:"
 OBJ_FEATS_DIR = 'assets/obj_features'
 
-
-def load_agent():
-    # build model
-    with open('model/cfg.yaml') as f:
-        cfg = yaml.safe_load(f)
-        cfg = OmegaConf.create(cfg)
-    agent = LeoAgentLLM(cfg)
-
-    # load checkpoint
-    if cfg.use_ckpt == 'hf':
-        ckpt_path = hf_hub_download(cfg.hf_ckpt_path[0], cfg.hf_ckpt_path[1])
-    else:
-        ckpt_path = cfg.local_ckpt_path
-    ckpt = torch.load(ckpt_path, map_location='cpu')
-    agent.load_state_dict(ckpt, strict=False)
-
-    agent.eval()
-    agent.to(DEVICE)
-    return agent
-
-agent = load_agent()
-
-
-def get_log_fname():
-    t = datetime.datetime.now()
-    fname = os.path.join(LOG_DIR, f'{t.year}-{t.month:02d}-{t.day:02d}.json')
-    return fname
+with open('cfg.yaml') as f:
+    cfg = yaml.safe_load(f)
+    cfg = OmegaConf.create(cfg)
+
+# build model
+agent = LeoAgentLLM(cfg)
+
+# load checkpoint
+if cfg.launch_mode == 'hf':
+    ckpt_path = hf_hub_download(cfg.hf_ckpt_path[0], cfg.hf_ckpt_path[1])
+else:
+    ckpt_path = cfg.local_ckpt_path
+ckpt = torch.load(ckpt_path, map_location='cpu')
+agent.load_state_dict(ckpt, strict=False)
+agent.eval()
+agent.to(DEVICE)
+
+os.makedirs(LOG_DIR, exist_ok=True)
+t = datetime.datetime.now()
+log_fname = os.path.join(LOG_DIR, f'{t.year}-{t.month:02d}-{t.day:02d}-{uuid4()}.json')
+
+if cfg.launch_mode == 'hf':
+    scheduler = CommitScheduler(
+        repo_id=cfg.hf_log_path,
+        repo_type='dataset',
+        folder_path=LOG_DIR,
+        path_in_repo=LOG_DIR,
+    )
 
 
 def change_scene(dropdown_scene: str):
@@ -139,17 +141,28 @@ def vote_response(
         'type': vote_type,
         'scene': dropdown_scene,
         'mode': dropdown_conversation_mode,
-        'dialogue': chatbot,
+        'dialogue': [chatbot[-1]] if 'Single-round' in dropdown_conversation_mode else chatbot,
     }
-    fname = get_log_fname()
-    if os.path.exists(fname):
-        with open(fname) as f:
-            logs = json.load(f)
-        logs.append(this_log)
+
+    if cfg.launch_mode == 'hf':
+        with scheduler.lock:   # use scheduler
+            if os.path.exists(log_fname):
+                with open(log_fname) as f:
+                    logs = json.load(f)
+                logs.append(this_log)
+            else:
+                logs = [this_log]
+            with open(log_fname, 'w') as f:
+                json.dump(logs, f, indent=2)
     else:
-        logs = [this_log]
-    with open(fname, 'w') as f:
-        json.dump(logs, f, indent=2)
+        if os.path.exists(log_fname):
+            with open(log_fname) as f:
+                logs = json.load(f)
+            logs.append(this_log)
+        else:
+            logs = [this_log]
+        with open(log_fname, 'w') as f:
+            json.dump(logs, f, indent=2)
 
 
 def upvote_response(