Spaces:
Running
Running
| import torch | |
| import sentencepiece | |
| from transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline | |
| import os | |
| import spacy | |
| import spacy_transformers | |
| import zipfile | |
| from collections import defaultdict | |
| class Models(): | |
| def __init__(self) -> None: | |
| self.load_trained_models() | |
| def load_trained_models(self): | |
| tokenizer = AutoTokenizer.from_pretrained("Jean-Baptiste/camembert-ner-with-dates",use_fast=False) | |
| model = AutoModelForTokenClassification.from_pretrained("Jean-Baptiste/camembert-ner-with-dates") | |
| self.ner = pipeline('ner', model=model, tokenizer=tokenizer, aggregation_strategy="simple") | |
| current_directory = os.path.dirname(os.path.realpath(__file__)) | |
| custom_ner_path = os.path.join(current_directory, 'spacy_model_v2/output/model-best') | |
| if not os.path.exists(custom_ner_path): | |
| with zipfile.ZipFile(r"models\prototype\spacy_model_v2.zip", 'r') as zip_ref: | |
| # Extract all contents in the current working directory | |
| zip_ref.extractall() | |
| self.custom_ner = spacy.load(custom_ner_path) | |
| def extract_ner(self, text): | |
| entities = self.ner(text) | |
| keys = ['DATE', 'ORG', 'LOC'] | |
| sort_dict = defaultdict(list) | |
| for entity in entities: | |
| if entity['score'] > 0.75: | |
| sort_dict[entity['entity_group']].append(entity['word']) | |
| filtered_dict = {key: value for key, value in sort_dict.items() if key in keys} | |
| filtered_dict = defaultdict(list, filtered_dict) | |
| return filtered_dict['DATE'], filtered_dict['ORG'], filtered_dict['LOC'] | |
| def get_ner(self, text, recover_text): | |
| dates, companies, locations = self.extract_ner(text) | |
| alternative_dates, alternative_companies, alternative_locations = self.extract_ner(recover_text) | |
| if dates == [] : | |
| dates = alternative_dates | |
| if companies == []: | |
| companies = alternative_companies | |
| if locations == []: | |
| locations = alternative_locations | |
| return dates, companies, locations | |
| def get_custom_ner(self, text): | |
| doc = self.custom_ner(text) | |
| entities = list(doc.ents) | |
| sort_dict = defaultdict(list) | |
| for entity in entities: | |
| sort_dict[entity.label_].append(entity.text) | |
| return sort_dict | |