def step(self, query: str) -> Tuple[torch.Tensor, float, bool, Dict]:
        """
        **Overview:**
            The step function of the language model environment.
        """
        # The history is not initialized. Create a new history.
        if self.history.tokens is None:
            query_tokens = self.tokenizer(query, return_tensors="pt").input_ids[0].to(self.current_device)
            self.history = TextHistory(query, query_tokens)
        # The history is already initialized. Append to the original history.
        else:
            query_tokens = self.tokenizer(query, return_tensors="pt").input_ids[0].to(self.current_device)
            self.history.append_segment(query, query_tokens)
        # Generate response.
        response_tokens = self.generate()
        # Calculate the reward function.
        rew = self.reward_fn(self.model, query_tokens, response_tokens)
        # Check whether the environment is finished.
        self.turn += 1
        self.history.completed = self.turn >= self.max_turns
        obs, mask = self.history.to_obs()
        return obs, rew, self.history.completed, {"mask": mask}

        **Overview:**
            The step function of the language model environment.
        
chapter8_large/lm_env.py