Spaces:
Build error
Build error
File size: 3,662 Bytes
171af38 baa262a 8ac889f 6489c85 d99c565 3a552ec 2d1b489 d99c565 3a552ec 171af38 3a552ec 8ac889f 3a552ec 171af38 01928f1 171af38 3a552ec 171af38 4770545 171af38 4770545 8310a9d 171af38 5d6f051 171af38 43b5bd6 171af38 8310a9d 171af38 efb6e7f 171af38 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | import json
import re
import logging
import numpy as np
import torch
from transformers import BertTokenizerFast, BertForTokenClassification
class Main:
def __init__(self):
self.MAX_LEN = 500
#self.EPOCHS = 6
#self.NUM_LABELS = 12
self.DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL_PATH = './models/bert-base-uncased'
self.STATE_DICT = torch.load('./models/model_e6.tar', map_location=self.DEVICE)
#tokenizer = BertTokenizerFast('../input/bert-base-uncased/vocab.txt', lowercase=True)
self.tokenizer = BertTokenizerFast('./models/bert-base-uncased/vocab.txt',lowercase=True)
self.model = BertForTokenClassification.from_pretrained(MODEL_PATH,state_dict=self.STATE_DICT['model_state_dict'], num_labels=12)
self.model.to(self.DEVICE);
self.tags_vals = ["UNKNOWN", "O", "Name", "Degree","Skills","College Name","Email Address","Designation","Companies worked at","Graduation Year","Years of Experience","Location"]
#self.tag2idx = {t: i for i, t in enumerate(tags_vals)}
self.idx2tag = {i:t for i, t in enumerate(self.tags_vals)}
#self.data = self.trim_entity_spans(self.convert_dataturks_to_spacy('./data/annotated_resumes.json'))
#self.model = MODEL
def process_resume(self,text):
max_len = self.MAX_LEN
tok = self.tokenizer.encode_plus(text, max_length=max_len, return_offsets_mapping=True)
curr_sent = dict()
padding_length = max_len - len(tok['input_ids'])
curr_sent['input_ids'] = tok['input_ids'] + ([0] * padding_length)
curr_sent['token_type_ids'] = tok['token_type_ids'] + ([0] * padding_length)
curr_sent['attention_mask'] = tok['attention_mask'] + ([0] * padding_length)
final_data = {
'input_ids': torch.tensor(curr_sent['input_ids'], dtype=torch.long),
'token_type_ids': torch.tensor(curr_sent['token_type_ids'], dtype=torch.long),
'attention_mask': torch.tensor(curr_sent['attention_mask'], dtype=torch.long),
'offset_mapping': tok['offset_mapping']
}
return final_data
def predict(self,test_resume):
return self._predict(self.model,test_resume)
def _predict(self,model,test_resume):
model.eval()
data = self.process_resume(test_resume)
input_ids, input_mask = data['input_ids'].unsqueeze(0), data['attention_mask'].unsqueeze(0)
labels = torch.tensor([1] * input_ids.size(0), dtype=torch.long).unsqueeze(0)
with torch.no_grad():
outputs = model(
input_ids,
token_type_ids=None,
attention_mask=input_mask,
labels=labels,
)
tmp_eval_loss, logits = outputs[:2]
logits = logits.cpu().detach().numpy()
label_ids = np.argmax(logits, axis=2)
entities = []
for label_id, offset in zip(label_ids[0], data['offset_mapping']):
curr_id = self.idx2tag[label_id]
curr_start = offset[0]
curr_end = offset[1]
if curr_id != 'O':
if len(entities) > 0 and entities[-1]['entity'] == curr_id and curr_start - entities[-1]['end'] in [0, 1]:
entities[-1]['end'] = curr_end
else:
entities.append({'entity': curr_id, 'start': curr_start, 'end':curr_end})
for ent in entities:
ent['text'] = test_resume[ent['start']:ent['end']]
return entities
|