Project: s41467-024-54639-7 | Function: MCTSEvaluator._get_reward (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:**         self._tokenizer = tokenizer
        self._bond_length_acceptability_cutoff = bond_length_acceptability_cutoff
        self._k = reward_k
        self._out_dir = out_dir
        self._num_valid = 0
        self._all_scores = []
        self._all_cifs = []

    def _postprocess(self, cif_str):
        # try to calculate the implied volume, to weed out very bad generations;
        #  an exception will be thrown if a value is missing, or the volume is nonsensical
        a = extract_numeric_property(cif_str, "_cell_length_a")
        b = extract_numeric_property(cif_str, "_cell_length_b")
        c = extract_numeric_property(cif_str, "_cell_length_c")
        alpha = extract_numeric_property(cif_str, "_cell_angle_alpha")
        beta = extract_numeric_property(cif_str, "_cell_angle_beta")
        gamma = extract_numeric_property(cif_str, "_cell_angle_gamma")
        get_unit_cell_volume(a, b, c, alpha, beta, gamma)

        # replace the symmetry operators with the correct operators
        space_group_symbol = extract_space_group_symbol(cif_str)
        if space_group_symbol is not None and space_group_symbol != "P 1":
            cif_str = replace_symmetry_operators(cif_str, space_group_symbol)

        # remove atom props
        cif_str = remove_atom_props_block(cif_str)

        return cif_str

    def _is_valid(self, generated_cif):
        if not is_formula_consistent(generated_cif):
            msg = "the generated CIF is inconsistent in terms of composition"
            return False, msg, None

        if not is_atom_site_multiplicity_consistent(generated_cif):
            msg = "the generated CIF is inconsistent in terms of atom site multiplicity"
            return False, msg, None

        bond_length_score = bond_length_reasonableness_score(generated_cif)
        if bond_length_score < self._bond_length_acceptability_cutoff:
            msg = f"unreasonable bond lengths detected " \
                  f"({(1 - bond_length_score) * 100:.0f}% of bond lengths were found to be unreasonable)"
            return False, msg, bond_length_score

        if not is_space_group_consistent(generated_cif):
            msg = "the generated CIF is inconsistent in terms of space group"
            return False, msg, None

        return True, "", None

    def _get_reward(self, score):
        """
        Returns a number between 0 and 1 representing the reward.
        Because we don't know the scale of the scores beforehand, we'll
        keep a running average of scores observed. A score which is close
        to the mean observed score will result in a reward close to 0.5.
        The k value determines how sensitive the reward is to deviations
        from the mean score.
        If higher scores are better, then provide a negative k value. If
        lower scores are better, provide a positive k value.
        """
        raise NotImplementedError("Function implementation has been masked")

    def _write_cif_to_file(self, cif, score, reward, id, iter_num):
        if self._out_dir is not None:

            if not os.path.exists(self._out_dir):
                os.makedirs(self._out_dir)

            cif_file = f"generated_{id}.cif"
            cif_fname = os.path.join(self._out_dir, cif_file)
            if not os.path.exists(cif_fname):
                print(f"writing CIF to: {cif_fname}")
                with open(cif_fname, "wt") as f:
                    f.write(cif)

                # create .csv to keep track of results
                csv_file = "results.csv"
                csv_fname = os.path.join(self._out_dir, csv_file)
                if not os.path.exists(csv_fname):
                    print(f"creating {csv_fname} as it does not exist...")
                    with open(csv_fname, "wt") as f:
                        f.write("file,iteration,score,reward\n")

                # update .csv
                with open(csv_fname, "a") as f:
                    f.write(f"{cif_file},{iter_num},{score},{reward}\n")

            else:
                print(f"CIF not written to file as it already exists: {cif_fname}")

    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


Complete the function implementation replacing "raise NotImplementedError".