INNER CODE UNIT · Python

_prepare_model

liyucheng09/Selective_Context · app/app.py:63

    def _prepare_model(self):
        if self.model_type == 'gpt2':
            if self.lang == 'zh':
                self.model = GPT2LMHeadModel.from_pretrained('uer/gpt2-chinese-cluecorpussmall')
                self.tokenizer = BertTokenizer.from_pretrained('uer/gpt2-chinese-cluecorpussmall')
            else:
                self.model = GPT2LMHeadModel.from_pretrained('gpt2')
                self.tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
            self.model.to(DEVICE)
            self.model.eval()

            print('model loaded')

            self.max_token_length = self.model.config.n_positions
            self.get_self_information = self._get_self_info_via_gpt2
    
    def get_self_information(self, text: str) -> Tuple[List[str], List[float]]:
        # it takes text as input, and return a list of words and a list of self-information scores

View source record →