def _save_model()

in src/autofill_model.py [0:0]

• 12 lines of code
• 3 McCabe index (conditional complexity)

    def _save_model(self, push_to_hub: bool = False):
        logger.info(f'Saving tokenizer and model to: {self.save_path}')
        label2id = self.dataset['train'].features["labels"]._str2int
        id2label = {v: k for k, v in label2id.items()}

        self.tokenizer.save_pretrained(self.save_path)
        self.trainer.save_model(self.save_path)

        # save with config
        config = AutoConfig.from_pretrained(self.save_path, label2id=label2id, id2label=id2label)
        model = AutoModelForSequenceClassification.from_pretrained(self.save_path, config=config)
        model.save_pretrained(self.save_path)

        if push_to_hub:
            # store on hugging face hub, requires login with huggingface-cli login
            logger.info('Saving to Hugging Face Hub...')
            self.trainer.push_to_hub(DEFAULT_HUGGING_FACE_PATH)