support streaming data, fix #284 #274 #268

Former-commit-id: 819cc1353599e5fa45658bc56dd0dbe4b258b197

src/llmtuner/tuner/ppo/trainer.py

@@ -2,21 +2,25 @@ import os
 import math
 import torch
 from tqdm import tqdm
-from typing import Callable, Dict, List, Optional
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional
 
-from transformers import Seq2SeqTrainingArguments, TrainerState, TrainerControl
+from transformers import TrainerState, TrainerControl
 from transformers.modeling_utils import PreTrainedModel
 
 from trl import PPOTrainer
 from trl.core import LengthSampler
 
-from llmtuner.extras.callbacks import LogCallback
 from llmtuner.extras.logging import get_logger
 from llmtuner.extras.misc import AverageMeter, get_logits_processor
-from llmtuner.hparams import FinetuningArguments
+
 from llmtuner.tuner.core.trainer import PeftTrainer
 from llmtuner.tuner.ppo.utils import cast_layernorm_dtype, replace_model
 
+if TYPE_CHECKING:
+    from transformers import Seq2SeqTrainingArguments
+    from llmtuner.extras.callbacks import LogCallback
+    from llmtuner.hparams import FinetuningArguments
+
 
 logger = get_logger(__name__)
 
@@ -27,9 +31,9 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
     """
     def __init__(
         self,
-        training_args: Seq2SeqTrainingArguments,
-        finetuning_args: FinetuningArguments,
-        callbacks: List[LogCallback],
+        training_args: "Seq2SeqTrainingArguments",
+        finetuning_args: "FinetuningArguments",
+        callbacks: List["LogCallback"],
         **kwargs
     ):
         PPOTrainer.__init__(self, **kwargs)

src/llmtuner/tuner/ppo/utils.py

@@ -1,11 +1,13 @@
 import torch
-from typing import Dict, List, Literal, Optional, Tuple
-from trl import AutoModelForCausalLMWithValueHead
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple
 
 from llmtuner.extras.constants import LAYERNORM_NAMES
 
+if TYPE_CHECKING:
+    from trl import AutoModelForCausalLMWithValueHead
 
-def replace_model(model: AutoModelForCausalLMWithValueHead, target: Literal["default", "reward"]) -> None:
+
+def replace_model(model: "AutoModelForCausalLMWithValueHead", target: Literal["default", "reward"]) -> None:
     if target == "reward": # save default head temporarily
         valuehead_state_dict = model.v_head.state_dict()
         setattr(model, "default_head_weight", valuehead_state_dict["summary.weight"])
@@ -19,10 +21,10 @@ def replace_model(model: AutoModelForCausalLMWithValueHead, target: Literal["def
 
 
 def cast_layernorm_dtype(
-    model: AutoModelForCausalLMWithValueHead,
+    model: "AutoModelForCausalLMWithValueHead",
     layer_norm_names: List[str] = LAYERNORM_NAMES,
     layer_norm_params: Optional[Dict[str, torch.Tensor]] = None
-) -> Tuple[AutoModelForCausalLMWithValueHead, Dict[str, torch.Tensor]]:
+) -> Tuple["AutoModelForCausalLMWithValueHead", Dict[str, torch.Tensor]]:
 
     layer_norm_state_dict = {}
 

src/llmtuner/tuner/ppo/workflow.py

@@ -2,26 +2,30 @@
 # https://github.com/lvwerra/trl/blob/main/examples/sentiment/scripts/gpt-neox-20b_peft/gpt-neo-20b_sentiment_peft.py
 
 import math
+from typing import TYPE_CHECKING
 from trl import PPOConfig
 from torch.optim import AdamW
 from typing import Optional, List
-from transformers import DataCollatorForSeq2Seq, Seq2SeqTrainingArguments, TrainerCallback
+from transformers import DataCollatorForSeq2Seq
 from transformers.optimization import get_scheduler
 
 from llmtuner.dsets import get_dataset, preprocess_dataset
 from llmtuner.extras.callbacks import LogCallback
 from llmtuner.extras.ploting import plot_loss
-from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
 from llmtuner.tuner.core import load_model_and_tokenizer
 from llmtuner.tuner.ppo.trainer import PPOPeftTrainer
 
+if TYPE_CHECKING:
+    from transformers import Seq2SeqTrainingArguments, TrainerCallback
+    from llmtuner.hparams import ModelArguments, DataArguments, FinetuningArguments
+
 
 def run_ppo(
-    model_args: ModelArguments,
-    data_args: DataArguments,
-    training_args: Seq2SeqTrainingArguments,
-    finetuning_args: FinetuningArguments,
-    callbacks: Optional[List[TrainerCallback]] = [LogCallback()]
+    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "Seq2SeqTrainingArguments",
+    finetuning_args: "FinetuningArguments",
+    callbacks: Optional[List["TrainerCallback"]] = [LogCallback()]
 ):
     dataset = get_dataset(model_args, data_args)
     model, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, training_args.do_train, stage="ppo")