[feat.] support encoder-decoder tuning and ChatGLM, Vicuna inference

examples/finetune.py

@@ -13,7 +13,7 @@
  foo = ClassFoo()
  bar = foo.FunctionBar()
 """
-
+import os
 import sys
 
 from transformers import HfArgumentParser
@@ -52,15 +52,31 @@ def main():
  pipeline_args=pipeline_args,
  )
  dataset = Dataset(data_args)
- model = AutoModel.get_model(model_args)
+
+ model = AutoModel.get_model(
+ model_args,
+ lang=data_args.lang,
+ forced_bos_token=data_args.forced_bos_token,
+ source_prefix = data_args.source_prefix,
+ streaming = data_args.streaming,
+ preprocessing_num_workers = data_args.preprocessing_num_workers,
+ overwrite_cache = data_args.overwrite_cache,
+ max_source_length = data_args.max_source_length,
+ max_target_length = data_args.max_target_length,
+ pad_to_max_length = data_args.pad_to_max_length
+ )
 
  # Tokenization and text grouping must be done in the main process
  with pipeline_args.main_process_first(desc="dataset map tokenization"):
  tokenized_dataset = model.tokenize(dataset)
- lm_dataset = finetuner.group_text(
- tokenized_dataset,
- model_max_length=model.get_max_length(),
- )
+ if model_args.arch_type == "encoder_decoder": 
+ # encoder-decoder model does not need group text
+ lm_dataset = tokenized_dataset
+ else:
+ lm_dataset = finetuner.group_text(
+ tokenized_dataset,
+ model_max_length=model.get_max_length(),
+ )
 
  # Finetuning
  tuned_model = finetuner.tune(model=model, lm_dataset=lm_dataset)

requirements.txt

@@ -6,6 +6,8 @@ wandb==0.14.0
 deepspeed==0.8.3
 trl @ git+https://github.com/lvwerra/trl.git#egg=trl-0.4.1
 sentencepiece
+icetk==0.0.7
+cpm_kernels==1.0.11
 transformers @ git+https://github.com/huggingface/transformers@c612628
 flask
-flask_cors
+flask_cors

scripts/run_chatbot_seq2seq.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+
+model=THUDM/chatglm-6b
+lora_args=""
+if [ $# -ge 1 ]; then
+ model=$1
+fi
+if [ $# -ge 2 ]; then
+ lora_args="--lora_model_path $2"
+fi
+
+CUDA_VISIBLE_DEVICES=0 \
+ deepspeed examples/chatbot.py \
+ --arch_type encoder_decoder \
+ --deepspeed configs/ds_config_chatbot.json \
+ --model_name_or_path ${model} \
+ ${lora_args}

src/lmflow/args.py

@@ -13,13 +13,14 @@
 """
 
 from dataclasses import dataclass, field
-from typing import Optional
+from pathlib import Path
+from typing import Optional, Union
 
 from transformers.utils.versions import require_version
-
+from transformers.generation.configuration_utils import GenerationConfig
 from transformers import (
  MODEL_FOR_CAUSAL_LM_MAPPING,
- TrainingArguments,
+ TrainingArguments 
 )
 
 MODEL_CONFIG_CLASSES = list(MODEL_FOR_CAUSAL_LM_MAPPING.keys())
@@ -99,6 +100,10 @@ class ModelArguments:
  default=None,
  metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)},
  )
+ arch_type: Optional[str] = field(
+ default="decoder_only",
+ metadata={"help": "The architecture type of the model. Currently supported decoder_only or encoder_decoder"}
+ )
  config_overrides: Optional[str] = field(
  default=None,
  metadata={
@@ -165,6 +170,15 @@ class ModelArguments:
  default=True,
  metadata={"help": "Whether use disk mapping when memory is not enough."}
  )
+ resize_position_embeddings: Optional[bool] = field(
+ default=None,
+ metadata={
+ "help": (
+ "Whether to automatically resize the position embeddings if `max_source_length` exceeds "
+ "the model's position embeddings."
+ )
+ },
+ )
 
  def __post_init__(self):
  if self.config_overrides is not None and (self.config_name is not None or self.model_name_or_path is not None):
@@ -225,6 +239,8 @@ class DatasetArguments:
  each parameter, such as a help message.
  """
 
+ lang: Optional[str] = field(default=None, metadata={"help": "Language id for summarization."})
+
  dataset_path: Optional[str] = field(
  default=None, metadata={"help": "The path of the dataset to use."}
  )
@@ -309,6 +325,83 @@ class DatasetArguments:
  default=None,
  metadata={"help": "Evaluation File Path"},
  )
+ max_source_length: Optional[int] = field(
+ default=1024,
+ metadata={
+ "help": (
+ "The maximum total input sequence length after tokenization. Sequences longer "
+ "than this will be truncated, sequences shorter will be padded."
+ )
+ },
+ )
+ max_target_length: Optional[int] = field(
+ default=128,
+ metadata={
+ "help": (
+ "The maximum total sequence length for target text after tokenization. Sequences longer "
+ "than this will be truncated, sequences shorter will be padded."
+ )
+ },
+ )
+ val_max_target_length: Optional[int] = field(
+ default=None,
+ metadata={
+ "help": (
+ "The maximum total sequence length for validation target text after tokenization. Sequences longer "
+ "than this will be truncated, sequences shorter will be padded. Will default to `max_target_length`."
+ "This argument is also used to override the ``max_length`` param of ``model.generate``, which is used "
+ "during ``evaluate`` and ``predict``."
+ )
+ },
+ )
+ pad_to_max_length: bool = field(
+ default=False,
+ metadata={
+ "help": (
+ "Whether to pad all samples to model maximum sentence length. "
+ "If False, will pad the samples dynamically when batching to the maximum length in the batch. More "
+ "efficient on GPU but very bad for TPU."
+ )
+ },
+ )
+ max_predict_samples: Optional[int] = field(
+ default=None,
+ metadata={
+ "help": (
+ "For debugging purposes or quicker training, truncate the number of prediction examples to this "
+ "value if set."
+ )
+ },
+ )
+ num_beams: Optional[int] = field(
+ default=None,
+ metadata={
+ "help": (
+ "Number of beams to use for evaluation. This argument will be passed to ``model.generate``, "
+ "which is used during ``evaluate`` and ``predict``."
+ )
+ },
+ )
+ ignore_pad_token_for_loss: bool = field(
+ default=True,
+ metadata={
+ "help": "Whether to ignore the tokens corresponding to padded labels in the loss computation or not."
+ },
+ )
+ source_prefix: Optional[str] = field(
+ default="", metadata={"help": "A prefix to add before every source text (useful for T5 models)."}
+ )
+
+ forced_bos_token: Optional[str] = field(
+ default=None,
+ metadata={
+ "help": (
+ "The token to force as the first generated token after the decoder_start_token_id."
+ "Useful for multilingual models like mBART where the first generated token"
+ "needs to be the target language token (Usually it is the target language token)"
+ )
+ },
+ )
 
  def __post_init__(self):
  if self.streaming:
@@ -330,7 +423,54 @@ class FinetunerArguments(TrainingArguments):
  """
  Adapt transformers.TrainingArguments
  """
- pass
+
+ """
+ Args:
+ sortish_sampler (`bool`, *optional*, defaults to `False`):
+ Whether to use a *sortish sampler* or not. Only possible if the underlying datasets are *Seq2SeqDataset*
+ for now but will become generally available in the near future.
+
+ It sorts the inputs according to lengths in order to minimize the padding size, with a bit of randomness
+ for the training set.
+ predict_with_generate (`bool`, *optional*, defaults to `False`):
+ Whether to use generate to calculate generative metrics (ROUGE, BLEU).
+ generation_max_length (`int`, *optional*):
+ The `max_length` to use on each evaluation loop when `predict_with_generate=True`. Will default to the
+ `max_length` value of the model configuration.
+ generation_num_beams (`int`, *optional*):
+ The `num_beams` to use on each evaluation loop when `predict_with_generate=True`. Will default to the
+ `num_beams` value of the model configuration.
+ """
+
+ sortish_sampler: bool = field(default=False, metadata={"help": "Whether to use SortishSampler or not."})
+ predict_with_generate: bool = field(
+ default=False, metadata={"help": "Whether to use generate to calculate generative metrics (ROUGE, BLEU)."}
+ )
+ generation_max_length: Optional[int] = field(
+ default=None,
+ metadata={
+ "help": (
+ "The `max_length` to use on each evaluation loop when `predict_with_generate=True`. Will default "
+ "to the `max_length` value of the model configuration."
+ )
+ },
+ )
+ generation_num_beams: Optional[int] = field(
+ default=None,
+ metadata={
+ "help": (
+ "The `num_beams` to use on each evaluation loop when `predict_with_generate=True`. Will default "
+ "to the `num_beams` value of the model configuration."
+ )
+ },
+ )
+ generation_config: Optional[Union[str, Path, GenerationConfig]] = field(
+ default=None,
+ metadata={
+ "help": "Model id, file path or url pointing to a GenerationConfig json file, to use during prediction."
+ },
+ )
+
 
 
 @dataclass

src/lmflow/datasets/dataset.py

@@ -169,7 +169,7 @@ def from_dict(self, dict_obj: dict, *args, **kwargs):
  return self
  else:
  raise NotImplementedError(
- f'Currently .from_dict is not supported for backend "{backend}"'
+ f'Currently .from_dict is not supported for backend "{self.backend}"'
  )
 
 
@@ -222,7 +222,7 @@ def to_dict(self):
  return dict_obj
  else:
  raise NotImplementedError(
- f'Current .to_dict is not supported for backend "{backend}"'
+ f'Current .to_dict is not supported for backend "{self.backend}"'
  )
 
 
@@ -251,7 +251,7 @@ def map(self, *args, **kwargs):
  else:
  # If the backend is not Hugging Face, raise a NotImplementedError
  raise NotImplementedError(
- f'Currently .map is not supported for backend "{backend}"'
+ f'Currently .map is not supported for backend "{self.backend}"'
  )
 
 

src/lmflow/models/auto_model.py

@@ -4,11 +4,16 @@
 """
 
 from lmflow.models.hf_decoder_model import HFDecoderModel
-
+from lmflow.models.hf_encoder_decoder_model import HFEncoderDecoderModel
 
 class AutoModel:
 
  @classmethod
  def get_model(self, model_args, *args, **kwargs):
  # TODO (add new models)
- return HFDecoderModel(model_args, *args, **kwargs)
+ if model_args.arch_type == "encoder_decoder":
+ return HFEncoderDecoderModel(model_args, *args, **kwargs)
+ elif model_args.arch_type == "decoder_only":
+ return HFDecoderModel(model_args, *args, **kwargs)
+ else:
+ raise NotImplementedError(f"Model type \"{model_args.arch_type}\" is not implemented.")

src/lmflow/models/encoder_decoder_model.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""A one-line summary of the module or program, terminated by a period.
+
+Leave one blank line. The rest of this docstring should contain an
+overall description of the module or program. Optionally, it may also
+contain a brief desription of exported classes and funcctions and/or usage
+examples.
+
+Typical usage example:
+
+ foo = ClassFoo()
+ bar = foo.FunctionBar()
+"""
+
+from lmflow.models.base_model import BaseModel
+
+
+class EncoderDecoderModel(BaseModel):
+
+ def __init__(self, *args, **kwargs):
+ pass

src/lmflow/models/hf_decoder_model.py

@@ -49,10 +49,8 @@
 from lmflow.models.decoder_model import DecoderModel
 from lmflow.models.interfaces.tunable import Tunable
 
-
 logger = logging.getLogger(__name__)
 
-
 class HFDecoderModel(DecoderModel, Tunable):
  r"""
  Initializes a HFDecoderModel instance.