INNER CODE UNIT · Python

_saved_pretrain

zjunlp/PromptKG · research/GenKGC/main.py:68

def _saved_pretrain(lit_model, tokenizer, path):
    lit_model.model.save_pretrained(path)
    tokenizer.save_pretrained(path)


def main():
    parser = _setup_parser()
    args = parser.parse_args()

    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    pl.seed_everything(args.seed)

    data_class = _import_class(f"data.{args.data_class}")
    model_class = _import_class(f"models.{args.model_class}")
    litmodel_class = _import_class(f"lit_models.{args.litmodel_class}")

    config = AutoConfig.from_pretrained(args.model_name_or_path)

View source record →