args.pretrained_model)
token_style = MODELS[args.pretrained_model][3]
# logs
model_save_path = args.weight_path
# Model
device = torch.device('cuda' if (
args.cuda and torch.cuda.is_available()) else 'cpu')
if args.use_crf:
deep_punctuation = DeepPunctuationCRF(args.pretrained_model,
freeze_bert=False,
lstm_dim=args.lstm_dim)
else:
deep_punctuation = DeepPunctuation(args.pretrained_model,
freeze_bert=False,
lstm_dim=args.lstm_dim)
deep_punctuation.to(device)
def inference():
deep_punctuation.load_state_dict(torch.load(model_save_path))
deep_punctuation.eval()
with open(args.in_file, 'r', encoding='utf-8') as f:
text = f.read()
text = re.sub(r"[,:\-–.!;?]", '', text)
words_original_case = text.split()
words = text.lower().split()
word_pos = 0
import re
import torch
import argparse
from model import DeepPunctuation, DeepPunctuationCRF
from config import *
pretrained_model = 'bert-base-multilingual-cased'
tokenizer = MODELS[pretrained_model][1].from_pretrained(pretrained_model)
token_style = MODELS[pretrained_model][3]
device = torch.device('cpu')
deep_punctuation = DeepPunctuation(pretrained_model, freeze_bert=False, lstm_dim=-1)
deep_punctuation.to(device)
def inference(model, text):
deep_punctuation.load_state_dict(torch.load(model, map_location=torch.device('cpu')), strict=False)
deep_punctuation.eval()
text = re.sub(r"[,:\-–.!;?]", '', text)
words_original_case = text.split()
words = text.lower().split()
word_pos = 0
sequence_len = 256
result = ""
decode_idx = 0
punctuation_map = {0: '', 1: ',', 2: '.', 3: '?'}
while word_pos < len(words):
'batch_size': args.batch_size,
'shuffle': True,
'num_workers': 1
}
# logs
os.makedirs(args.save_path, exist_ok=True)
model_save_path = os.path.join(args.save_path, 'weights.pt')
st_model_save_path = os.path.join(args.save_path, 'st_weights.pt')
log_path = os.path.join(args.save_path, args.name + '_logs.txt')
# Model
device = torch.device('cuda' if (args.cuda and torch.cuda.is_available()) else 'cpu')
deep_punctuation = DeepPunctuation(args.pretrained_model, freeze_bert=args.freeze_bert, lstm_dim=args.lstm_dim)
deep_punctuation.to(device)
criterion = nn.CrossEntropyLoss()
mse = nn.MSELoss()
# optimizer = torch.optim.Adam(deep_punctuation.parameters(), lr=args.lr, weight_decay=args.decay)
deep_punctuation.load_state_dict(torch.load(model_save_path, map_location=torch.device(device)))
with open(args.student) as f:
cfg = json.load(f)
studet_deep = Student(cfg)
studet_deep.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(studet_deep.parameters(), lr=args.lr, weight_decay=args.decay)
def validate(data_loader):
if args.yttm == 'false':
tokenizer = MODELS[args.pretrained_model][1].from_pretrained(
args.pretrained_model)
else:
tokenizer = YTTM(args.yttm)
# tokenizer = MODELS[args.pretrained_model][1].from_pretrained(args.pretrained_model)
token_style = MODELS[args.pretrained_model][3]
# logs
#model_save_path = args.weight_path
st_model_save_path = args.weight_path
# Model
device = torch.device('cuda' if (
args.cuda and torch.cuda.is_available()) else 'cpu')
if args.yttm == 'false':
deep_punctuation = DeepPunctuation(args.pretrained_model,
lstm_dim=args.lstm_dim)
else:
if args.pqrnn:
deep_punctuation = PQRNN()
else:
deep_punctuation = BiLSTM_CNN_CRF(4, tokenizer.vocab_size, 512)
# deep_punctuation.to(device)
with open(args.student) as f:
cfg = json.load(f)
#deep_punctuation = Student(cfg)
deep_punctuation.to(device)
newModuleList = nn.ModuleList()
# Now iterate over all layers, only keepign only the relevant layers.
for i in range(0, len(num_layers_to_keep)):
newModuleList.append(oldModuleList[i])
# create a copy of the model, modify it with the new list, and return
copyOfModel = copy.deepcopy(model)
copyOfModel.bert_layer.encoder.layer = newModuleList
return copyOfModel
if __name__ == '__main__':
device = torch.device('cpu')
deep_punctuation = DeepPunctuation(args.pretrained_model, freeze_bert=False, lstm_dim=-1)
model_save_path = os.path.join(args.save_path, 'dp_bert_wts.pt')
new_model_save_path = os.path.join(args.save_path, 'pruned_weights.pt')
deep_punctuation.to(device)
deep_punctuation.load_state_dict(torch.load(model_save_path, map_location=torch.device(device)))
# seq = torch.nn.Sequential(*(list(deep_punctuation.children())))
# print(len(seq))
# # print(deep_punctuation)
# bert_module = seq[0]
# seq_bert = torch.nn.Sequential(*(list(bert_module.encoder.children())))
# print(len(bert_module))
ls = [0,6,13]
new_model = deleteEncodingLayers(deep_punctuation, ls)
print(new_model)