def __init__(self, model: GPT2LMHeadModel, tokenizer: GPT2Tokenizer, reward_fn: Callable,
                 max_turns: int = 4, generation_kwargs: Optional[Dict] = None):
        """
        **Overview:**
            Initialize the TextEnvironment.

        **Arguments:**
            - model: The model to use for generation.
            - tokenizer: The tokenizer to use for generation.
            - reward_fn: A callable function that takes a string and returns a reward.
            - max_turns: The maximum number of turns to allow.
            - generation_kwargs: A dictionary of keyword arguments to pass to the model's generate method.
        """
        # Initialize the arguments.
        self.model = model
        self.tokenizer = tokenizer
        self.reward_fn = reward_fn
        self.max_turns = max_turns

        # Prepare the arguments for text generation.
        if generation_kwargs is None:
            self.generation_kwargs = dict()
        else:
            self.generation_kwargs = generation_kwargs

        # Count the times of ``self.step()``
        self.turn = 0
        # Preserve the history of interactions.
        self.history = TextHistory("", None)
        # Determine the device of running the model (cpu or cuda).
        self.current_device = self.model.device

        # Define the action-space, reward-space and observation-space.
        # The action space is a sentence (string type).
        self._action_space = gym.spaces.Text(max_length=1024)
        # In this demo, we use the negative perplexity as reward, whose range is (-inf, 0).
        self._reward_space = gym.spaces.Box(-float('inf'), 0)
        # The observation is the history query and response, whose shape is 1024.
        # If the total length of history < 1024, it will be padded. Detailed implementation is shown in ``TextHistory.to_obs``.
        # For each element of the observation, the value range is [0, vcab_size).
        self._observation_space = gym.spaces.Box(0, tokenizer.vocab_size, [1024])

        **Overview:**
            Initialize the TextEnvironment.

        **Arguments:**
            - model: The model to use for generation.
            - tokenizer: The tokenizer to use for generation.
            - reward_fn: A callable function that takes a string and returns a reward.
            - max_turns: The maximum number of turns to allow.
            - generation_kwargs: A dictionary of keyword arguments to pass to the model's generate method.
        
chapter8_large/lm_env.py