updated loading code

Signed-off-by: ftgreat <[email protected]>

examples/Aquila/Aquila-chat/aquila_chat.py

@@ -66,26 +66,29 @@
       flush=True)
 
 checkpoints = env_args.pre_load_dir
-
 model_name = env_args.model_name
 
 print('*' * 20, "model_name", model_name, flush=True)
-
 cache_dir = os.path.join(checkpoints, model_name)
 print('*' * 20, "cache_dir", cache_dir)
-tokenizer = Tokenizer.from_pretrained(model_name, cache_dir=cache_dir)
-print('*' * 20, "tokenizer", tokenizer)
 
 # avoid sync loading models in case of Mem OOM
 if env_args.bmt_async_load:
     import time
     time.sleep(10 * 60 * (trainer.local_rank % 4))
 
 config_file = os.path.join(cache_dir, 'config.json')
-from flagai.model.aquila_model import AQUILAModel
-
-model = AQUILAModel.init_from_json(config_file=config_file)
-# print('*'*20, "model", model)
+# from flagai.model.aquila_model import AQUILAModel
+
+loader = AutoLoader("lm",
+                    model_dir='./checkpoints_in/',
+                    model_name=model_name,
+                    use_cache=False,
+                    fp16=True)
+model = loader.get_model()
+tokenizer = loader.get_tokenizer()
+# print('*' * 20, "tokenizer", tokenizer)
+print('*'*20, "model", model)
 
 #lora
 if env_args.lora:

examples/Aquila/Aquila-chat/hostfile

@@ -1 +1 @@
-192.168.21.4 slots=4
+192.168.21.2 slots=4

flagai/env_trainer_v1.py

@@ -698,7 +698,7 @@ def do_train(self,
                                     save_optim=self.save_optim,
                                     save_dir=self.save_dir,
                                     save_rng=self.save_rng,
-                                    iteration_in_epoch=iteration_)
+                                    iteration_in_epoch=99999999999)
 
         # Evaluation #todo add train_args
         if ((self.epochs == 0) or (self.eval_interval and