text: #!/usr/bin/env python3
  # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
  
  import itertools
  import math
  from typing import Any, Dict, Iterable, List
  
  import torch
  from pytext.common.constants import DatasetFieldName, DFColumn, VocabMeta
  from pytext.config.field_config import FeatureConfig, WordLabelConfig
  from pytext.data.featurizer import InputRecord
  from pytext.fields import TextFeatureField
  from pytext.utils import cuda
  from torchtext import data as textdata
  
  from .data_handler import BatchIterator, DataHandler
  
  
  class BPTTLanguageModelDataHandler(DataHandler):
      """
      `BPTTLanguageModelDataHandler` treats data as a single document, concatenating
      all tokens together. BPTTIterator arranges the dataset into columns of
      batch size and subdivides the source data into chunks of length bptt_len.
      It enables hidden state of ith batch carried over to (i+1)th batch.
  
      Args:
          bptt_len (int) : Input sequence length to backpropagate to.
      """
  
      class Config(DataHandler.Config):
          """
          Configuration class for `BPTTLanguageModelDataHandler`.
  
          Attributes:
              columns_to_read (List[str]): List containing the names of the
                  columns to read from the data files.
              bptt_len (int): Input sequence length to backpropagate to.
          """
  
          columns_to_read: List[str] = [DFColumn.UTTERANCE]
          bptt_len: int = 35
  
      def __init__(self, bptt_len: int, **kwargs) -> None:
          super().__init__(**kwargs)
          self.bptt_len = bptt_len
  
      @classmethod
      def from_config(
          cls,
          config: Config,
          feature_config: FeatureConfig,
          label_config: WordLabelConfig,
          **kwargs
      ):
          """
          Factory method to construct an instance of `BPTTLanguageModelDataHandler`
          from the module's config object and feature config object.
  
          Args:
              config (LanguageModelDataHandler.Config): Configuration object
                  specifying all the parameters of `BPTTLanguageModelDataHandler`.
              feature_config (FeatureConfig): Configuration object specifying all
                  the parameters of all input features.
  
          Returns:
              type: An instance of `BPTTLanguageModelDataHandler`.
          """
          # For language modeling the only input is a collection of utterances.
          # The input and the labels are created by the LangaugeModelDataHandler.
          # The input at time step t+1 becomes a label for the input at time step t.
          columns = config.columns_to_read
          bptt_len = config.bptt_len
          if bptt_len <= 0:
              raise TypeError("BPTT Sequence length cannot be 0 or less.")
          features = {
              # the name must be text because it's hardcoded in torchtext BPTT iterator
              "text": TextFeatureField(
                  eos_token=VocabMeta.EOS_TOKEN, include_lengths=False
              )
          }
          return cls(
              bptt_len=bptt_len,
              raw_columns=columns,
              features=features,
              labels={},
              extra_fields={},
              train_path=config.train_path,
              eval_path=config.eval_path,
              test_path=config.test_path,
              train_batch_size=config.train_batch_size,
              eval_batch_size=config.eval_batch_size,
              test_batch_size=config.test_batch_size,
              pass_index=False,
              **kwargs
          )
  
      def init_feature_metadata(
          self,
          train_data: textdata.Dataset,
          eval_data: textdata.Dataset,
          test_data: textdata.Dataset,
      ):
          """
          Prepares the metadata for the language model features.
          """
          super().init_feature_metadata(train_data, eval_data, test_data)
          # workaround the hardcoded torchtext name
          self.metadata.features[DatasetFieldName.TEXT_FIELD] = self.metadata.features[
              "text"
          ]
  
      def init_target_metadata(
          self,
          train_data: textdata.Dataset,
          eval_data: textdata.Dataset,
          test_data: textdata.Dataset,
      ):
          """
          Prepares the metadata for the language model target.
          """
          self.metadata.target = self.metadata.features[DatasetFieldName.TEXT_FIELD]
  
      def preprocess(self, data: Iterable[Dict[str, Any]]):
          tokens = []
          for row in data:
              tokens.extend(self.preprocess_row(row))
          return [{"text": tokens}]
  
      def preprocess_row(self, row_data: Dict[str, Any]) -> List[str]:
          """
          Preprocess steps for a single input row.
  
          Args:
              row_data (Dict[str, Any]): Dict representing the input row and
                  columns.
  
          Returns:
              List[str]: List of tokens.
          """
          return self.featurizer.featurize(
              InputRecord(raw_text=row_data[DFColumn.UTTERANCE])
          ).tokens
  
      def _get_train_iter(
          self, shard_dataset: textdata.Dataset, batch_size: int, world_size: int = 1
      ) -> BatchIterator:
          # Compute the per-worker batch size
          assert (
              batch_size >= world_size
          ), "batch size needs to be >= the distributed world size"
          batch_size = batch_size // world_size
  
          return BatchIterator(
              textdata.BPTTIterator(
                  shard_dataset,
                  batch_size=batch_size,
                  bptt_len=self.bptt_len,
                  device="cuda:{}".format(torch.cuda.curren_device())
                  if cuda.CUDA_ENABLED
                  else "cpu",
                  sort_within_batch=True,
                  repeat=False,
                  sort_key=self.sort_key,
              ),
              self._postprocess_batch,
              num_batches=math.ceil(len(shard_dataset) / float(batch_size)),
          )
  
      def _train_input_from_batch(self, batch):
  
          return (
              # (bsx x seq_len)
              batch.text.t().contiguous(),
              torch.Tensor([batch.text.size(0)] * batch.text.size(1)).type_as(batch.text),
          )
  
      def _target_from_batch(self, batch):
          return (batch.target.t().contiguous(),)
  
      def get_test_iter(self, file_path: str, batch_size: int) -> BatchIterator:
          """
          Get test data iterator from test data file.
  
          Args:
              file_path (str): Path to test data file.
              batch_size (int): Batch size
  
          Returns:
              BatchIterator: An instance of BatchIterator to iterate over the
                  supplied test data file.
          """
          test_data = self.gen_dataset_from_path(file_path)
          return BatchIterator(
              textdata.BPTTIterator(
                  test_data,
                  batch_size=batch_size,
                  bptt_len=self.bptt_len,
                  device="cuda:{}".format(torch.cuda.curren_device())
                  if cuda.CUDA_ENABLED
                  else "cpu",
                  sort=True,
                  repeat=False,
                  train=False,
                  sort_key=self.sort_key,
              ),
              self._postprocess_batch,
          )
  