Project: s41467-024-54639-7 | Function: MCTSLanguageModel.rollout (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** The generation of plausible crystal structures is often the first step in predicting the structure and properties of a material from its chemical composition. However, most current methods for crystal structure prediction are computationally expensive, slowing the pace of innovation. Seeding structure prediction algorithms with quality generated candidates can overcome a major bottleneck. Here, we introduce CrystaLLM, a methodology for the versatile generation of crystal structures, based on the autoregressive large language modeling (LLM) of the Crystallographic Information File (CIF) format. Trained on millions of CIF files, CrystaLLM focuses on modeling crystal structures through text. CrystaLLM can produce plausible crystal structures for a wide range of inorganic compounds unseen in training, as demonstrated by ab initio simulations. Our approach challenges conventional representations of crystals, and demonstrates the potential of LLMs for learning effective models of crystal chemistry, which will lead to accelerated discovery and innovation in materials science.

**KEY TERMS:** {
  "AIRSS": "Examples are the AIRSS ^6,7 and USPEX ^8 approaches",
  "CIF": "Even so, the model is capable of reliably generating correct CIF syntax and physically plausible crystal structures for many classes of inorganic compounds",
  "USPEX": "Examples are the AIRSS ^6,7 and USPEX ^8 approaches",
  "Recent": "Recent studies have demonstrated that LLMs trained on sequences of board game play (e",
  "These": "These types of approaches can also be enhanced to increase the structural diversity of generated materials, by allowing partial substitutions and adjusting substitution probabilities ^68",
  "The Large Language Model": "The Large Language Model (LLM) approach, backed by the Transformer architecture ^19 , is behind state-of-the-art performance on natural language processing tasks"
}

**CONTEXT CODE:** 
    def __call__(self, token_sequence, iter_num):
        cif = self._tokenizer.decode(token_sequence)

        try:
            cif = self._postprocess(cif)
            valid, msg, bond_length_score = self._is_valid(cif)
            if not valid:
                print(f"CIF invalid: {msg}")
                if bond_length_score is not None:
                    return -(1 - bond_length_score)
                else:
                    return -1.0
        except Exception as e:
            print(f"exception while post-processing and validating: {e}")
            print(traceback.format_exc())
            return -1.0

        self._num_valid += 1

        try:
            print("invoking external scorer...")
            score = self._scorer.score(cif)
            print(f"external scorer returned score: {score}")
        except Exception as e:
            print(f"exception while scoring: {e}")
            print(traceback.format_exc())
            return -1.0

        if math.isnan(score):
            print(f"reward cannot be computed as score is nan")
            return -1.0

        reward = self._get_reward(score)
        print(f"computed reward: {reward}")

        self._write_cif_to_file(cif, score, reward, self._num_valid, iter_num)

        return reward


class MCTSLanguageModel:
    def __init__(self, model: GPT, config: GPTConfig, child_ids: List[int], device: str, temperature: float):
        self._model = model
        self._model.eval()
        self._config = config
        self._child_ids = child_ids
        self._device = device
        self._temperature = temperature

    def rollout(self, rollout_state: List[int], width: int, max_depth: int, newline_id: int) -> List[int]:
        """Performs MCTS rollout by generating tokens from the current node using the language model"""
        raise NotImplementedError("Function implementation has been masked")

    def top_n_vocab_with_weights(self, n: int, token_sequence: List[int]) -> Tuple[List[int], List[float]]:
        idx = (torch.tensor(token_sequence, dtype=torch.long, device=self._device)[None, ...])

        # if the sequence context is growing too long we must crop it at block_size
        idx_cond = idx if idx.size(1) <= self._config.block_size else idx[:, -self._config.block_size:]
        # forward the model to get the logits for the index in the sequence
        logits, _ = self._model(idx_cond)
        # pluck the logits at the final step and scale by desired temperature
        logits = logits[:, -1, :] / self._temperature

        dist = torch.distributions.categorical.Categorical(logits=logits)

        tokens_and_log_probs = []
        for child_id in self._child_ids:
            new_token_sequence = list(token_sequence)
            new_token_sequence.append(child_id)
            log_prob = dist.log_prob(torch.tensor(child_id, device=self._device).view(1)).item()
            tokens_and_log_probs.append((child_id, log_prob))

        top_n = list(reversed(sorted(tokens_and_log_probs, key=lambda k: k[1])))[:n]

        top_n_child_ids = [t[0] for t in top_n]
        top_n_weights = self._normalize([t[1] for t in top_n])

        return top_n_child_ids, top_n_weights

    @staticmethod
    def _normalize(log_probs: List[float]) -> List[float]:
        probs = [math.exp(lp) for lp in log_probs]
        prob_factor = 1 / sum(probs)
        return [prob_factor * p for p in probs]


class ContextSensitiveTreeBuilder:
    def __init__(
        self,
        tokenizer: CIFTokenizer,
        top_child_weight_cutoff: float = 0.99,
        n_space_groups: int = 0,
        bypass_only_child: bool = False,
    ):
        self._tok = tokenizer
        self._top_child_weight_cutoff = top_child_weight_cutoff
        self._n_space_groups = n_space_groups
        self._bypass_only_child = bypass_only_child

    def get_child_ids_and_weights(
        self,
        state: List[int],
        top_n_child_ids: List[int],
        top_n_weights: List[float],
        lm: MCTSLanguageModel,
        width: int,
        newline_id: int,
    ) -> Tuple[Union[List[int], List[List[int]]], List[float]]:

        if len(state) > 1 and state[-2:] == [self._tok.token_to_id["_symmetry_space_group_name_H-M"], self._tok.token_to_id[" "]] and self._n_space_groups > 0:
            return lm.top_n_vocab_with_weights(self._n_space_groups, state)

        top_child_id = top_n_child_ids[0]
        top_child_weight = top_n_weights[0]
        if top_child_weight > self._top_child_weight_cutoff:

            if self._bypass_only_child:
                only_children = []
                while top_child_weight > self._top_child_weight_cutoff:
                    only_children.append(top_child_id)
                    new_state = state + only_children
                    if MCTSNode.is_complete(new_state, newline_id):
                        return [only_children], [1.]
                    top_n_child_ids, top_n_weights = lm.top_n_vocab_with_weights(width, new_state)
                    top_child_id = top_n_child_ids[0]
                    top_child_weight = top_n_weights[0]

                top_n_extended_child_ids = [only_children + [child_id] for child_id in top_n_child_ids]
                return top_n_extended_child_ids, top_n_weights

            return [top_child_id], [1.]

        return top_n_child_ids, top_n_weights


class MCTSNode:
    def __init__(
        self,
        state: List[int],
        language_model: MCTSLanguageModel,
        width: int,
        max_depth: int,
        newline_id: int,
        parent: "MCTSNode" = None,
        tree_builder: ContextSensitiveTreeBuilder = None,
    ):
        self.state = state
        self._newline_id = newline_id
        self._lm = language_model


Complete the function implementation replacing "raise NotImplementedError".