update models

app.py

@@ -1,7 +1,5 @@
 from transformers import AutoModel, AutoTokenizer
 from copy import deepcopy
-import os
-import ipdb
 import gradio as gr
 import mdtex2html
 from model.openlamm import LAMMPEFTModel

model/__init__.py

@@ -1,10 +1 @@
-# from .agent import DeepSpeedAgent
 from .openlamm import LAMMPEFTModel
-
-
-# def load_model(args):
-#     agent_name = args['models'][args['model']]['agent_name']
-#     model_name = args['models'][args['model']]['model_name']
-#     model = globals()[model_name](**args)
-#     agent = globals()[agent_name](model, args)
-#     return agent

model/agent.py

@@ -1,76 +0,0 @@
-from header import *
-from torch.utils.tensorboard import SummaryWriter
-
-
-class DeepSpeedAgent:
-    
-    def __init__(self, model, args):
-        super(DeepSpeedAgent, self).__init__()
-        self.args = args
-        self.model = model
-        self.writer = SummaryWriter(args['log_path'])
-        if args['stage'] == 2:
-            self.load_stage_1_parameters(args["delta_ckpt_path"])
-            print(f'[!] load stage 1 checkpoint from {args["delta_ckpt_path"]}')
-
-        # load config parameters of deepspeed
-        ds_params = json.load(open(self.args['ds_config_path']))
-        ds_params['scheduler']['params']['total_num_steps'] = self.args['total_steps']
-        ds_params['scheduler']['params']['warmup_num_steps'] = max(10, int(self.args['total_steps'] * self.args['warmup_rate']))
-        self.ds_engine, self.optimizer, _ , _ = deepspeed.initialize(
-            model=self.model, 
-            model_parameters=self.model.parameters(),
-            config_params=ds_params, 
-            dist_init_required=True,
-            args=types.SimpleNamespace(**args)
-        )
-
-    @torch.no_grad()
-    def predict(self, batch):
-        self.model.eval()
-        string = self.model.generate_one_sample(batch)
-        return string
-
-    def train_model(self, batch, current_step=0, pbar=None):
-        self.ds_engine.module.train()
-        loss, mle_acc = self.ds_engine(batch)
-
-        self.ds_engine.backward(loss)
-        self.ds_engine.step()
-        pbar.set_description(f'[!] loss: {round(loss.item(), 4)}; token_acc: {round(mle_acc*100, 2)}')
-        pbar.update(1)
-        if self.args['local_rank'] == 0 and self.args['log_path'] and current_step % self.args['logging_step'] == 0:
-            elapsed = pbar.format_dict['elapsed']
-            rate = pbar.format_dict['rate']
-            remaining = (pbar.total - pbar.n) / rate if rate and pbar.total else 0
-            remaining = str(datetime.timedelta(seconds=remaining))
-            self.writer.add_scalar('train/loss', loss.item(), current_step)
-            self.writer.add_scalar('train/token_acc', mle_acc*100, current_step)
-            logging.info(f'[!] progress: {round(pbar.n/pbar.total, 5)}; remaining time: {remaining}; loss: {round(loss.item(), 4)}; token_acc: {round(mle_acc*100, 2)}')
-            
-        mle_acc *= 100
-        return mle_acc
-    
-    def save_model(self, path, current_step):
-        # only save trainable model parameters
-        param_grad_dic = {
-            k: v.requires_grad for (k, v) in self.ds_engine.module.named_parameters()
-        }
-        state_dict = self.ds_engine.module.state_dict()
-        checkpoint = OrderedDict()
-        for k, v in self.ds_engine.module.named_parameters():
-            if v.requires_grad:
-                checkpoint[k] = v
-        if current_step <= 0:
-            torch.save(checkpoint, f'{path}/pytorch_model.pt')
-        else:
-            torch.save(checkpoint, f'{path}/pytorch_model_ep{current_step}.pt')
-        # save tokenizer
-        self.model.llama_tokenizer.save_pretrained(path)
-        # save configuration
-        self.model.llama_model.config.save_pretrained(path)
-        print(f'[!] save model into {path}')
-
-    def load_stage_1_parameters(self, path):
-        delta_ckpt = torch.load(path, map_location=torch.device('cpu'))
-        self.model.load_state_dict(delta_ckpt, strict=False)

model/openlamm.py

@@ -21,7 +21,7 @@ from .CLIP import load as load_clip
 from .PROCESS import data
 from .modeling_llama import LlamaForCausalLM
 from .utils.pcl_utils import MEAN_COLOR_RGB, RandomCuboid, random_sampling
-import .conversations
+from .conversations import conversation_dict, default_conversation
 
 ImageFile.LOAD_TRUNCATED_IMAGES = True
 
@@ -139,9 +139,9 @@ def make_prompt_start(system_header=False, vision_type='image', task_type='norma
     PROMPT_START = f'### Human: {VISION_TAGS["sov"][vision_type]}'
     if system_header:
         if task_type == 'normal':
-            return f"{conversations.default_conversation.system}\n\n" + PROMPT_START
+            return f"{default_conversation.system}\n\n" + PROMPT_START
         else:
-            return [f"{conversations.conversation_dict[task]}\n\n" + PROMPT_START for task in task_type]
+            return [f"{conversation_dict[task]}\n\n" + PROMPT_START for task in task_type]
     else:
         return PROMPT_START