update code structure

README.md

@@ -166,7 +166,8 @@ sys.path.append("src")
 from src import load_pretrained, ModelArguments
 model_args = ModelArguments(checkpoint_dir=path_to_checkpoint)
 model, tokenizer = load_pretrained(model_args)
-model = model.half().cuda()
+model = model.cuda()
+model.eval()
 # model.generate, model.chat()...
 ```
 

README_zh.md

@@ -171,7 +171,8 @@ sys.path.append("src")
 from src import load_pretrained, ModelArguments
 model_args = ModelArguments(checkpoint_dir=path_to_checkpoint)
 model, tokenizer = load_pretrained(model_args)
-model = model.half().cuda()
+model = model.cuda()
+model.eval()
 # model.generate, model.chat()...
 ```
 

examples/evaluate.sh

@@ -4,8 +4,9 @@ CUDA_VISIBLE_DEVICES=0 python ../src/finetune.py \
     --do_eval \
     --dataset alpaca_gpt4_zh \
     --dataset_dir ../data \
-    --output_dir output_eval \
+    --checkpoint_dir path_to_checkpoint \
+    --output_dir path_to_eval_result \
     --overwrite_cache \
     --per_device_eval_batch_size 8 \
-    --max_samples 20 \
+    --max_samples 50 \
     --predict_with_generate

examples/finetune.sh

@@ -5,14 +5,13 @@ CUDA_VISIBLE_DEVICES=0 python ../src/finetune.py \
     --dataset alpaca_gpt4_zh \
     --dataset_dir ../data \
     --finetuning_type lora \
-    --output_dir output_finetune \
+    --output_dir path_to_checkpoint \
     --overwrite_cache \
     --per_device_train_batch_size 4 \
     --gradient_accumulation_steps 4 \
     --lr_scheduler_type cosine \
     --logging_steps 10 \
     --save_steps 1000 \
-    --max_samples 10000 \
     --learning_rate 5e-5 \
     --num_train_epochs 1.0 \
     --fp16

src/finetune.py

@@ -9,9 +9,9 @@
     prepare_data,
     preprocess_data,
     plot_loss,
-    DataCollatorForChatGLM,
+    Seq2SeqDataCollatorForChatGLM,
     ComputeMetrics,
-    TrainerForChatGLM
+    Seq2SeqTrainerForChatGLM
 )
 
 
@@ -20,9 +20,9 @@ def main():
     # Prepare pretrained model and dataset
     model_args, data_args, training_args, finetuning_args = prepare_args()
     dataset = prepare_data(model_args, data_args)
-    model, tokenizer = load_pretrained(model_args, finetuning_args, is_trainable=training_args.do_train)
+    model, tokenizer = load_pretrained(model_args, training_args, finetuning_args, is_trainable=training_args.do_train)
     dataset = preprocess_data(dataset, tokenizer, data_args, training_args)
-    data_collator = DataCollatorForChatGLM(
+    data_collator = Seq2SeqDataCollatorForChatGLM(
         tokenizer=tokenizer,
         model=model,
         ignore_pad_token_for_loss=data_args.ignore_pad_token_for_loss,
@@ -36,7 +36,7 @@ def main():
                 data_args.num_beams is not None else training_args.generation_num_beams
 
     # Initialize our Trainer
-    trainer = TrainerForChatGLM(
+    trainer = Seq2SeqTrainerForChatGLM(
         finetuning_args=finetuning_args,
         model=model,
         args=training_args,
@@ -67,14 +67,12 @@ def main():
 
     # Evaluation
     if training_args.do_eval:
-        model = model.half() # don't use `--fp16` argument at evaluation
         metrics = trainer.evaluate(metric_key_prefix="eval", **gen_kwargs)
         trainer.log_metrics("eval", metrics)
         trainer.save_metrics("eval", metrics)
 
     # Predict
     if training_args.do_predict:
-        model = model.half()
         predict_results = trainer.predict(dataset, metric_key_prefix="predict", **gen_kwargs)
         trainer.log_metrics("predict", predict_results.metrics)
         trainer.save_metrics("predict", predict_results.metrics)

src/infer.py

@@ -36,7 +36,8 @@ def main():
     parser = HfArgumentParser(ModelArguments)
     model_args, = parser.parse_args_into_dataclasses()
     model, tokenizer = load_pretrained(model_args)
-    model = model.half().cuda()
+    model = model.cuda()
+    model.eval()
 
     history = []
     print(welcome)

src/utils/__init__.py

@@ -2,10 +2,13 @@
     load_pretrained,
     prepare_args,
     prepare_data,
-    preprocess_data,
-    DataCollatorForChatGLM,
+    preprocess_data
+)
+
+from .seq2seq import (
+    Seq2SeqDataCollatorForChatGLM,
     ComputeMetrics,
-    TrainerForChatGLM
+    Seq2SeqTrainerForChatGLM
 )
 
 from .config import ModelArguments