implement do_predict

README.md

@@ -127,9 +127,21 @@ CUDA_VISIBLE_DEVICES=0 python src/finetune.py \
     --do_eval \
     --dataset alpaca_gpt4_zh \
     --checkpoint_dir path_to_checkpoint \
-    --output_dir eval \
+    --output_dir path_to_eval_result \
     --per_device_eval_batch_size 8 \
-    --max_eval_samples 50 \
+    --max_samples 50 \
+    --predict_with_generate
+```
+
+### Predict
+```bash
+CUDA_VISIBLE_DEVICES=0 python src/finetune.py \
+    --do_predict \
+    --dataset alpaca_gpt4_zh \
+    --checkpoint_dir path_to_checkpoint \
+    --output_dir path_to_predict_result \
+    --per_device_eval_batch_size 8 \
+    --max_samples 50 \
     --predict_with_generate
 ```
 
@@ -149,7 +161,7 @@ CUDA_VISIBLE_DEVICES=0 python src/web_demo.py \
 ### Deploy the Fine-tuned Model
 ```python
 from .src import load_pretrained, ModelArguments
-model_args = ModelArguments(checkpoint_dir=path_to_checkpoint_dir)
+model_args = ModelArguments(checkpoint_dir=path_to_checkpoint)
 model, tokenizer = load_pretrained(model_args)
 model = model.half().cuda()
 # model.generate, model.chat()...

README_zh.md

@@ -129,31 +129,43 @@ CUDA_VISIBLE_DEVICES=0 python src/finetune.py \
     --do_eval \
     --dataset alpaca_gpt4_zh \
     --checkpoint_dir path_to_checkpoint \
-    --output_dir eval \
+    --output_dir path_to_eval_result \
     --per_device_eval_batch_size 8 \
-    --max_eval_samples 50 \
+    --max_samples 50 \
     --predict_with_generate
 ```
 
-### 效果测试
+### 模型预测
+```bash
+CUDA_VISIBLE_DEVICES=0 python src/finetune.py \
+    --do_predict \
+    --dataset alpaca_gpt4_zh \
+    --checkpoint_dir path_to_checkpoint \
+    --output_dir path_to_predict_result \
+    --per_device_eval_batch_size 8 \
+    --max_samples 50 \
+    --predict_with_generate
+```
+
+### 命令行测试
 
 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/infer.py \
     --checkpoint_dir path_to_checkpoint
 ```
 
-### 网页版测试
+### 浏览器测试
 
 ```bash
 CUDA_VISIBLE_DEVICES=0 python src/web_demo.py \
     --checkpoint_dir path_to_checkpoint
 ```
 
-### 部署微调模型
+### 模型部署
 
 ```python
 from .src import load_pretrained, ModelArguments
-model_args = ModelArguments(checkpoint_dir=path_to_checkpoint_dir)
+model_args = ModelArguments(checkpoint_dir=path_to_checkpoint)
 model, tokenizer = load_pretrained(model_args)
 model = model.half().cuda()
 # model.generate, model.chat()...

examples/evaluate.sh

@@ -7,5 +7,5 @@ CUDA_VISIBLE_DEVICES=0 python ../src/finetune.py \
     --output_dir output_eval \
     --overwrite_cache \
     --per_device_eval_batch_size 8 \
-    --max_eval_samples 20 \
+    --max_samples 20 \
     --predict_with_generate

examples/finetune.sh

@@ -12,7 +12,7 @@ CUDA_VISIBLE_DEVICES=0 python ../src/finetune.py \
     --lr_scheduler_type cosine \
     --logging_steps 10 \
     --save_steps 1000 \
-    --max_train_samples 10000 \
+    --max_samples 10000 \
     --learning_rate 5e-5 \
     --num_train_epochs 1.0 \
     --fp16

src/finetune.py

@@ -19,10 +19,15 @@ def main():
 
     # Prepare pretrained model and dataset
     model_args, data_args, training_args, finetuning_args = prepare_args()
-    dataset = prepare_data(model_args, data_args, training_args)
+    dataset = prepare_data(model_args, data_args)
     model, tokenizer = load_pretrained(model_args, finetuning_args, is_trainable=training_args.do_train)
     dataset = preprocess_data(dataset, tokenizer, data_args, training_args)
-    data_collator = DataCollatorForChatGLM(tokenizer, model, data_args.ignore_pad_token_for_loss, training_args.do_eval)
+    data_collator = DataCollatorForChatGLM(
+        tokenizer=tokenizer,
+        model=model,
+        ignore_pad_token_for_loss=data_args.ignore_pad_token_for_loss,
+        inference_mode=(not training_args.do_train)
+    )
 
     # Override the decoding parameters of Trainer
     training_args.generation_max_length = training_args.generation_max_length if \
@@ -42,6 +47,14 @@ def main():
         compute_metrics=ComputeMetrics(tokenizer) if training_args.predict_with_generate else None
     )
 
+    # Keyword arguments for `model.generate`
+    gen_kwargs = {
+        "do_sample": True,
+        "top_p": 0.7,
+        "max_length": 768,
+        "temperature": 0.95
+    }
+
     # Training
     if training_args.do_train:
         train_result = trainer.train()
@@ -55,10 +68,18 @@ 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", do_sample=True, top_p=0.7, max_length=768, temperature=0.95)
+        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)
+        trainer.save_predictions(predict_results, tokenizer)
+
 
 def _mp_fn(index):
     # For xla_spawn (TPUs)

src/utils/common.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import json
 import hashlib
 import logging
 import torch
@@ -18,7 +19,7 @@
     DataCollatorForSeq2Seq,
     set_seed
 )
-from transformers.trainer import TRAINING_ARGS_NAME
+from transformers.trainer import PredictionOutput, TRAINING_ARGS_NAME
 from transformers.deepspeed import is_deepspeed_zero3_enabled
 from transformers.modeling_utils import PreTrainedModel
 from transformers.tokenization_utils import PreTrainedTokenizer
@@ -47,7 +48,8 @@
     prepare_model_for_training,
     merge_lora_weights,
     IGNORE_INDEX,
-    FINETUNING_ARGS_NAME
+    FINETUNING_ARGS_NAME,
+    PREDICTION_FILE_NAME
 )
 
 
@@ -146,11 +148,13 @@ def load_pretrained(
                 checkpoints_to_merge, lastest_checkpoint = model_args.checkpoint_dir[:-1], model_args.checkpoint_dir[-1]
             else:
                 checkpoints_to_merge = model_args.checkpoint_dir
-            merge_lora_weights(model, checkpoints_to_merge)
+            checkpoint_merged = merge_lora_weights(model, checkpoints_to_merge)
             if lastest_checkpoint is not None: # resume lora training
                 model = PeftModel.from_pretrained(model, lastest_checkpoint, is_trainable=is_trainable)
                 if not is_trainable: # merge the lastest LoRA weights to evaluate the model
                     model.merge_and_unload()
+                    checkpoint_merged += 1
+            logger.info("Merged {} model checkpoint(s).".format(checkpoint_merged))
         if lastest_checkpoint is None: # create new lora weights
             lora_config = LoraConfig(
                 task_type=TaskType.CAUSAL_LM,
@@ -180,8 +184,8 @@ def prepare_args() -> Tuple[ModelArguments, DataTrainingArguments, Seq2SeqTraini
         model_args, data_args, training_args, finetuning_args = parser.parse_args_into_dataclasses()
 
     # Check arguments
-    if training_args.do_train and training_args.do_eval:
-        raise ValueError("We don't support training and evaluation simultaneously.")
+    if int(training_args.do_train) + int(training_args.do_eval) + int(training_args.do_predict) != 1:
+        raise ValueError("We must perform single operation from do_train, do_eval and do_predict.")
     training_args.optim = "adamw_torch" if training_args.optim == "adamw_hf" else training_args.optim # suppress warning
 
     if model_args.quantization_bit is not None: # perform FP16 checking or GPU checking
@@ -225,8 +229,7 @@ def prepare_args() -> Tuple[ModelArguments, DataTrainingArguments, Seq2SeqTraini
 
 def prepare_data(
         model_args: ModelArguments,
-        data_args: DataTrainingArguments,
-        training_args: Seq2SeqTrainingArguments
+        data_args: DataTrainingArguments
 ) -> Dataset:
     # Load and verify dataset
     def checksum(file_path, hash):
@@ -236,7 +239,7 @@ def checksum(file_path, hash):
         if sha1 != hash:
             logger.warning("Checksum failed for {}. It may vary depending on the platform.".format(file_path))
 
-    max_samples = data_args.max_train_samples if training_args.do_train else data_args.max_eval_samples
+    max_samples = data_args.max_samples
     all_datasets = [] # support multiple datasets
 
     for dataset_info in data_args.dataset_list:
@@ -533,3 +536,22 @@ def prediction_step(
             labels = None
 
         return loss, generated_tokens, labels
+
+    def save_predictions(
+            self,
+            predict_results: PredictionOutput,
+            tokenizer: PreTrainedTokenizer
+    ) -> None: # custom behavior
+        if self.is_world_process_zero():
+            if self.args.predict_with_generate:
+                predictions = tokenizer.batch_decode(predict_results.predictions, skip_special_tokens=True)
+                predictions = [pred.strip() for pred in predictions]
+                labels = tokenizer.batch_decode(predict_results.label_ids, skip_special_tokens=True)
+                labels = [label.strip() for label in labels]
+                output_prediction_file = os.path.join(self.args.output_dir, PREDICTION_FILE_NAME)
+                logger.info(f"Saving prediction results to {output_prediction_file}")
+                with open(output_prediction_file, "w", encoding="utf-8") as writer:
+                    res = []
+                    for pred, label in zip(predictions, labels):
+                        res.append(json.dumps({"label": label, "predict": pred}, ensure_ascii=False))
+                    writer.write("\n".join(res))

src/utils/config.py

@@ -56,10 +56,6 @@ class ModelArguments:
         default=False,
         metadata={"help": "Will use the token generated when running `huggingface-cli login`."}
     )
-    resize_position_embeddings: Optional[bool] = field(
-        default=False,
-        metadata={"help": "Whether to resize the position embeddings if `max_source_length` exceeds or not."}
-    )
     quantization_bit: Optional[int] = field(
         default=None,
         metadata={"help": "The number of bits to quantize the model."}
@@ -107,13 +103,9 @@ class DataTrainingArguments:
         default=512,
         metadata={"help": "The maximum total output sequence length after tokenization."}
     )
-    max_train_samples: Optional[int] = field(
-        default=None,
-        metadata={"help": "For debugging purposes, truncate the number of training examples for each dataset."}
-    )
-    max_eval_samples: Optional[int] = field(
+    max_samples: Optional[int] = field(
         default=None,
-        metadata={"help": "For debugging purposes, truncate the number of evaluation examples for each dataset."}
+        metadata={"help": "For debugging purposes, truncate the number of examples for each dataset."}
     )
     num_beams: Optional[int] = field(
         default=None,

src/utils/other.py

@@ -90,7 +90,7 @@ def forward(self, x):
 
 # This function merges lora weights from multiple checkpoints
 # Inspired by: https://github.com/huggingface/peft/blob/34027fe813756897767b9a6f19ae7f1c4c7b418c/src/peft/tuners/lora.py#L451
-def merge_lora_weights(model: PreTrainedModel, checkpoints_to_merge: List[str]) -> None:
+def merge_lora_weights(model: PreTrainedModel, checkpoints_to_merge: List[str]) -> int:
     checkpoint_merged = 0
     for checkpoint_dir in checkpoints_to_merge:
         adapter_config = json.load(open(os.path.join(checkpoint_dir, CONFIG_NAME), "r"))
@@ -114,7 +114,7 @@ def merge_lora_weights(model: PreTrainedModel, checkpoints_to_merge: List[str])
                 param.data += weight_to_merge.to(param.device) * scaling
                 is_merged = True
         checkpoint_merged = checkpoint_merged + 1 if is_merged else checkpoint_merged
-    logger.info("Merged {} model checkpoint(s).".format(checkpoint_merged))
+    return checkpoint_merged
 
 
 def plot_loss(training_args: Seq2SeqTrainingArguments) -> None:
@@ -136,3 +136,4 @@ def plot_loss(training_args: Seq2SeqTrainingArguments) -> None:
 
 IGNORE_INDEX = -100
 FINETUNING_ARGS_NAME = "finetuning_args.bin"
+PREDICTION_FILE_NAME = "generated_predictions.txt"