chatAI / save_model.py
gaur3009's picture
Upload save_model.py
6c54217 verified
raw
history blame contribute delete
751 Bytes
#!/usr/bin/env python
# coding: utf-8
# In[1]:
import torch
from transformers import BertTokenizer, BertModel, GPT2LMHeadModel, GPT2Tokenizer
# Load pre-trained models and tokenizers
bert_model_name = 'bert-base-uncased'
bert_tokenizer = BertTokenizer.from_pretrained(bert_model_name)
bert_model = BertModel.from_pretrained(bert_model_name)
gpt2_model_name = 'gpt2'
gpt2_tokenizer = GPT2Tokenizer.from_pretrained(gpt2_model_name)
gpt2_model = GPT2LMHeadModel.from_pretrained(gpt2_model_name)
# Save BERT model and tokenizer
bert_model.save_pretrained("bert_model")
bert_tokenizer.save_pretrained("bert_model")
# Save GPT-2 model and tokenizer
gpt2_model.save_pretrained("gpt2_model")
gpt2_tokenizer.save_pretrained("gpt2_model")
# In[ ]: