# Buggy code
def __init__(self, args, tgt_dict):
    super().__init__(args, tgt_dict)

    self.unit_lm = getattr(args, "unit_lm", False)

    self.lexicon = load_words(args.lexicon) if args.lexicon else None
    self.idx_to_wrd = {}

    checkpoint = torch.load(args.kenlm_model, map_location="cpu")

    if "cfg" in checkpoint and checkpoint["cfg"] is not None:
        lm_args = checkpoint["cfg"]
    else:
        lm_args = convert_namespace_to_omegaconf(checkpoint["args"])

    with open_dict(lm_args.task):
        lm_args.task.data = osp.dirname(args.kenlm_model)

    task = tasks.setup_task(lm_args.task)
    model = task.build_model(lm_args.model)
    model.load_state_dict(checkpoint["model"], strict=False)

    self.trie = Trie(self.vocab_size, self.silence)

    self.word_dict = task.dictionary
    self.unk_word = self.word_dict.unk()
    self.lm = FairseqLM(self.word_dict, model)

    if self.lexicon:
        start_state = self.lm.start(False)
        for i, (word, spellings) in enumerate(self.lexicon.items()):
            if self.unit_lm:
                word_idx = i
                self.idx_to_wrd[i] = word
                score = 0
            else:
                word_idx = self.word_dict.index(word)
                _, score = self.lm.score(start_state, word_idx, no_cache=True)

            for spelling in spellings:
                spelling_idxs = [tgt_dict.index(token) for token in spelling]
                assert tgt_dict.unk() not in spelling_idxs, (
                    f"{spelling} {spelling_idxs}"
                )
                self.trie.insert(spelling_idxs, word_idx, score)
        self.trie.smear(SmearingMode.MAX)

        self.decoder_opts = LexiconDecoderOptions(
            beam_size=args.beam,
            beam_size_token=int(getattr(args, "beam_size_token", len(tgt_dict))),
            beam_threshold=args.beam_threshold,
            lm_weight=args.lm_weight,
            word_score=args.word_score,
            unk_score=args.unk_weight,
            sil_score=args.sil_weight,
            log_add=False,
            criterion_type=self.criterion_type,
        )

        self.decoder = LexiconDecoder(
            self.decoder_opts,
            self.trie,
            self.lm,
            self.silence,
            self.blank,
            self.unk_word,
            self.asg_transitions,
            self.unit_lm,
        )
    else:
        assert args.unit_lm, (
            "lexicon free decoding can only be done with a unit language model"
        )
        from flashlight.lib.text.decoder import (
            LexiconFreeDecoder,
            LexiconFreeDecoderOptions,
        )

        d = {w: [[w]] for w in tgt_dict.symbols}
        self.word_dict = create_word_dict(d)
        self.lm = KenLM(args.kenlm_model, self.word_dict)
        self.decoder_opts = LexiconFreeDecoderOptions(
            beam_size=args.beam,
            beam_size_token=int(getattr(args, "beam_size_token", len(tgt_dict))),
            beam_threshold=args.beam_threshold,
            lm_weight=args.lm_weight,
            sil_score=args.sil_weight,
            log_add=False,
            criterion_type=self.criterion_type,
        )
        self.decoder = LexiconFreeDecoder(
            self.decoder_opts, self.lm, self.silence, self.blank, []
        )


# Traceback (KeyError)
Traceback (most recent call last):
File "attention_layers.py", line 80, in <module>
tokens = roberta.encode(' '.join(list(s)))
File "/home/jmorton/software/fairseq/fairseq/models/roberta/hub_interface.py", line 57, in encode
bpe_sentence = '
 ' + self.bpe.encode(sentence) + ' 
'
File "/home/jmorton/software/fairseq/fairseq/data/encoders/gpt2_bpe.py", line 40, in encode
return ' '.join(map(str, self.bpe.encode(x)))
File "/home/jmorton/software/fairseq/fairseq/data/encoders/gpt2_bpe_utils.py", line 110, in encode
bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
File "/home/jmorton/software/fairseq/fairseq/data/encoders/gpt2_bpe_utils.py", line 110, in <genexpr>
bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
KeyError: 'Ġ'

# Fixed code
def __init__(self, args, tgt_dict):
    super().__init__(args, tgt_dict)

    self.unit_lm = getattr(args, "unit_lm", False)

    self.lexicon = load_words(args.lexicon) if args.lexicon else None
    self.idx_to_wrd = {}

    checkpoint = torch.load(args.kenlm_model, map_location="cpu")

    if "cfg" in checkpoint and checkpoint["cfg"] is not None:
        lm_args = checkpoint["cfg"]
    else:
        lm_args = convert_namespace_to_omegaconf(checkpoint["args"])

    with open_dict(lm_args.task):
        lm_args.task.data = osp.dirname(args.kenlm_model)

    task = tasks.setup_task(lm_args.task)
    model = task.build_model(lm_args.model)
    model.load_state_dict(checkpoint["model"], strict=False)

    self.trie = Trie(self.vocab_size, self.silence)

    self.word_dict = task.dictionary
    self.unk_word = self.word_dict.unk()
    self.lm = FairseqLM(self.word_dict, model)

    if self.lexicon:
        start_state = self.lm.start(False)
        for i, (word, spellings) in enumerate(self.lexicon.items()):
            if self.unit_lm:
                word_idx = i
                self.idx_to_wrd[i] = word
                score = 0
            else:
                word_idx = self.word_dict.index(word)
                _, score = self.lm.score(start_state, word_idx, no_cache=True)

            for spelling in spellings:
                spelling_idxs = [tgt_dict.index(token) for token in spelling]
                assert tgt_dict.unk() not in spelling_idxs, (
                    f"{spelling} {spelling_idxs}"
                )
                self.trie.insert(spelling_idxs, word_idx, score)
        self.trie.smear(SmearingMode.MAX)

        self.decoder_opts = LexiconDecoderOptions(
            beam_size=args.beam,
            beam_size_token=int(getattr(args, "beam_size_token", len(tgt_dict))),
            beam_threshold=args.beam_threshold,
            lm_weight=args.lm_weight,
            word_score=args.word_score,
            unk_score=args.unk_weight,
            sil_score=args.sil_weight,
            log_add=False,
            criterion_type=self.criterion_type,
        )

        self.decoder = LexiconDecoder(
            self.decoder_opts,
            self.trie,
            self.lm,
            self.silence,
            self.blank,
            self.unk_word,
            [],
            self.unit_lm,
        )
    else:
        assert args.unit_lm, (
            "lexicon free decoding can only be done with a unit language model"
        )
        from flashlight.lib.text.decoder import (
            LexiconFreeDecoder,
            LexiconFreeDecoderOptions,
        )

        d = {w: [[w]] for w in tgt_dict.symbols}
        self.word_dict = create_word_dict(d)
        self.lm = KenLM(args.kenlm_model, self.word_dict)
        self.decoder_opts = LexiconFreeDecoderOptions(
            beam_size=args.beam,
            beam_size_token=int(getattr(args, "beam_size_token", len(tgt_dict))),
            beam_threshold=args.beam_threshold,
            lm_weight=args.lm_weight,
            sil_score=args.sil_weight,
            log_add=False,
            criterion_type=self.criterion_type,
        )
        self.decoder = LexiconFreeDecoder(
            self.decoder_opts, self.lm, self.silence, self.blank, []
        )