resume-parser / Models.py
Sybghat's picture
Initial commit
324a80e
raw
history blame
2.27 kB
from transformers import AutoTokenizer, AutoModelForTokenClassification, AutoModelForSequenceClassification
from transformers import pipeline
from flair.data import Sentence
from flair.models import SequenceTagger
import pickle
class Models:
def pickle_it(self, obj, file_name):
with open(f'{file_name}.pickle', 'wb') as f:
pickle.dump(obj, f)
def unpickle_it(self, file_name):
with open(f'{file_name}.pickle', 'rb') as f:
return pickle.load(f)
def load_trained_models(self, pickle=False):
#NER (dates)
tokenizer = AutoTokenizer.from_pretrained("Jean-Baptiste/camembert-ner-with-dates")
model = AutoModelForTokenClassification.from_pretrained("Jean-Baptiste/camembert-ner-with-dates")
self.ner_dates = pipeline('ner', model=model, tokenizer=tokenizer, aggregation_strategy="simple")
#Zero Shot Classification
# self.zero_shot_classifier = pipeline("zero-shot-classification", model='facebook/bart-large-mnli')
self.zero_shot_classifier = pipeline("zero-shot-classification", model='valhalla/distilbart-mnli-12-6')
# Ner
tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
model = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER")
self.ner = pipeline('ner', model=model, tokenizer=tokenizer, grouped_entities=True)
# Pos Tagging
self.tagger = SequenceTagger.load("flair/pos-english-fast")
if pickle:
self.pickle_models()
return self.ner, self.ner_dates, self.zero_shot_classifier, self.tagger
def pickle_models(self):
self.pickle_it(self.ner, "ner")
self.pickle_it(self.zero_shot_classifier, "zero_shot_classifier_6")
self.pickle_it(self.ner_dates, "ner_dates")
self.pickle_it(self.tagger, "pos_tagger_fast")
def load_pickled_models(self):
ner_dates = self.unpickle_it('ner_dates')
ner = self.unpickle_it('ner')
zero_shot_classifier = self.unpickle_it('zero_shot_classifier_6')
tagger = self.unpickle_it("pos_tagger_fast")
return ner_dates, ner, zero_shot_classifier, tagger
def get_flair_sentence(self, sent):
return Sentence(sent)