docstr: Trains the model.
docstr_tok: ["Trains","the","model","."]
code: def train(self, output_dir: str) -> None:
      """
      Trains the model.
      """
      if isinstance(self._task, TrainingTaskForPPO):
          if not all(
              x in self.training_args_kwargs
              for x in [
                  "length_sampler_kwargs",
                  "generation_kwargs",
                  "reward_model",
              ]
          ):
              raise ValueError(
                  "To train a PPO model, you need to specify the following arguments via `trainer.update_config`: length_sampler_kwargs, generation_kwargs, reward_model."
              )
  
      from transformers import TrainingArguments
  
      # check required path argument
      self.training_args_kwargs["output_dir"] = output_dir
  
      self.init_model(new=True)
  
      if isinstance(self._task, TrainingTaskForSFT):
          self._trainer = self.trainer_cls(
              self.trainer_model,
              args=TrainingArguments(**self.training_args_kwargs),
              train_dataset=self._train_dataset,
              eval_dataset=self._eval_dataset,
              dataset_text_field="text",
              tokenizer=self.trainer_tokenizer,
              **self.trainer_kwargs,
          )
  
      elif isinstance(self._task, TrainingTaskForRM):
  
          def preprocess_function(examples) -> Dict[str, List]:
              new_examples = {
                  "input_ids_chosen": [],
                  "attention_mask_chosen": [],
                  "input_ids_rejected": [],
                  "attention_mask_rejected": [],
              }
              for chosen, rejected in zip(
                  examples["chosen"], examples["rejected"]
              ):
                  tokenized_j = self.trainer_tokenizer(
                      chosen, truncation=True
                  )
                  tokenized_k = self.trainer_tokenizer(
                      rejected, truncation=True
                  )
  
                  new_examples["input_ids_chosen"].append(
                      tokenized_j["input_ids"]
                  )
                  new_examples["attention_mask_chosen"].append(
                      tokenized_j["attention_mask"]
                  )
                  new_examples["input_ids_rejected"].append(
                      tokenized_k["input_ids"]
                  )
                  new_examples["attention_mask_rejected"].append(
                      tokenized_k["attention_mask"]
                  )
  
              return new_examples
  
          train_dataset = self._train_dataset.map(
              preprocess_function, batched=True
          )
          eval_dataset = None
          if self._eval_dataset:
              eval_dataset = self._eval_dataset.map(
                  preprocess_function, batched=True
              )
  
          self._trainer = self.trainer_cls(
              self.trainer_model,
              args=TrainingArguments(**self.training_args_kwargs),
              train_dataset=train_dataset,
              eval_dataset=eval_dataset,
              tokenizer=self.trainer_tokenizer,
              **self.trainer_kwargs,
          )
      elif isinstance(self._task, TrainingTaskForPPO):
          from datasets import concatenate_datasets
  
          dataset = concatenate_datasets(
              [
                  x
                  for x in [self._train_dataset, self._eval_dataset]
                  if x is not None
              ]
          )
  
          def tokenize(sample):
              sample["input_ids"] = self.trainer_tokenizer.encode(
                  sample["query"], truncation=True
              )
              return sample
  
          def data_collator(data):
              return dict(
                  (key, [d[key] for d in data]) for key in data[0]
              )
  
          def remove_truncated(sample):
              return (
                  len(sample) < self.trainer_tokenizer.model_max_length
              )
  
          dataset = dataset.map(tokenize, batched=False)
          size_before = len(dataset)
          dataset = dataset.filter(remove_truncated, batched=False)
          size_after = len(dataset)
          if size_after != size_before:
              self._logger.info(
                  f"Removed {size_before - size_after} samples ({1 - (size_after / size_before):%}), "
                  "as these samples were longer than the maximum model length even before the generation."
              )
          dataset.set_format(type="torch")
          self._trainer = self.trainer_cls(
              model=self.trainer_model,
              ref_model=self.trainer_ref_model,
              tokenizer=self.trainer_tokenizer,
              dataset=dataset,
              data_collator=data_collator,
              **self.trainer_kwargs,
          )
      elif isinstance(self._task, TrainingTaskForDPO):
          self._trainer = self.trainer_cls(
              model=self.trainer_model,
              ref_model=self.trainer_ref_model,
              args=TrainingArguments(**self.training_args_kwargs),
              train_dataset=self._train_dataset,
              eval_dataset=self._eval_dataset,
              tokenizer=self.trainer_tokenizer,
              **self.trainer_kwargs,
          )
  
      #  train
      if isinstance(self._task, TrainingTaskForPPO):
          import torch
          from tqdm import tqdm
          from trl.core import LengthSampler
  
          output_length_sampler = LengthSampler(
              **self.training_args_kwargs["length_sampler_kwargs"]
          )
          generation_kwargs = self.training_args_kwargs[
              "generation_kwargs"
          ]
          generation_kwargs[
              "pad_token_id"
          ] = self.trainer_tokenizer.eos_token_id
          reward_model = self.training_args_kwargs["reward_model"]
  
          for batch in tqdm(self._trainer.dataloader):
              query_tensors = batch["input_ids"]
  
              #### Get response from SFT
              response_tensors = self._trainer.generate(
                  query_tensors,
                  return_prompt=False,
                  length_sampler=output_length_sampler,
                  **generation_kwargs,
              )
              batch["response"] = self.trainer_tokenizer.batch_decode(
                  response_tensors, skip_special_tokens=True
              )
  
              #### Compute rewards scores
              texts = [
                  q + r
                  for q, r in zip(batch["query"], batch["response"])
              ]
              pipe_outputs = reward_model(
                  texts, top_k=None, truncation=True
              )
              rewards = [
                  torch.tensor(output[-1]["score"] * len(output))
                  for output in pipe_outputs
              ]
  
              #### Run PPO step
              stats = self._trainer.step(
                  query_tensors, response_tensors, rewards
              )
              self._trainer.log_stats(stats, batch, rewards)
      else:
          self._trainer.train()
          if self._trainer.eval_dataset:
              self._metrics = self._trainer.evaluate()
              self._logger.info(self._metrics)
          else:
              self._metrics = None
  
      self.save(output_dir)
  
file_path: argilla/src/argilla/client/feedback/training/frameworks/trl.py
imports: import logging
  from typing import TYPE_CHECKING, Dict, List, Optional, Union
  from argilla.client.feedback.training.base import ArgillaTrainerSkeleton
  from argilla.client.feedback.training.schemas.base import TrainingTaskForDPO, TrainingTaskForPPO, TrainingTaskForRM, TrainingTaskForSFT
  from argilla.training.utils import filter_allowed_args
  from argilla.utils.dependency import require_dependencies, requires_dependencies
label_desc: Use image node_id 6 for calling the ArgillaTRLTrainer obj's underlying member method code with example usage: obj.train(output_dir) and returns: new_examples, sample, dict, unknown