src/rime/models/zero_shot/bayes_lm.py [41:50]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        assert self.tokenizer.padding_side == 'right', "expect right padding"
        if model_name == 'gpt2':
            self.tokenizer.pad_token = self.tokenizer.eos_token

        if model_name.startswith('bert'):
            self.model = BertForMaskedLM.from_pretrained(model_name)
        else:  # gpt2
            self.model = AutoModelForCausalLM.from_pretrained(model_name)
        self.model.eval()  # eval mode
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



src/rime/models/zero_shot/item_knn.py [36:45]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
            assert self.tokenizer.padding_side == 'right', "expect right padding"
            if model_name == 'gpt2':
                self.tokenizer.pad_token = self.tokenizer.eos_token

            if model_name.startswith('bert'):
                self.model = BertForMaskedLM.from_pretrained(model_name)
            else:  # gpt2
                self.model = AutoModelForCausalLM.from_pretrained(model_name)
            self.model.eval()  # eval mode
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -