def __init__()

in source/embed.py [0:0]

• 21 lines of code
• 7 McCabe index (conditional complexity)

    def __init__(self, model_path, max_sentences=None, max_tokens=None, cpu=False, fp16=False, verbose=False,
                 sort_kind='quicksort'):
        self.use_cuda = torch.cuda.is_available() and not cpu
        self.max_sentences = max_sentences
        self.max_tokens = max_tokens
        if self.max_tokens is None and self.max_sentences is None:
            self.max_sentences = 1

        state_dict = torch.load(model_path)
        self.encoder = Encoder(**state_dict['params'])
        self.encoder.load_state_dict(state_dict['model'])
        self.dictionary = state_dict['dictionary']
        self.pad_index = self.dictionary['<pad>']
        self.eos_index = self.dictionary['</s>']
        self.unk_index = self.dictionary['<unk>']
        if fp16:
            self.encoder.half()
        if self.use_cuda:
            if verbose:
                print(' - transfer encoder to GPU')
            self.encoder.cuda()
        self.sort_kind = sort_kind