| import os
|
| from transformers import AutoTokenizer, AutoModelForTokenClassification
|
| import torch
|
|
|
| MODEL_PATH = os.path.abspath(
|
| os.path.join(os.path.dirname(__file__), "..", "models", "ner_model")
|
| )
|
|
|
|
|
| if not os.path.isdir(MODEL_PATH):
|
| MODEL_PATH = "4deffo/tamil-ner-model"
|
|
|
| tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
|
| model = AutoModelForTokenClassification.from_pretrained(MODEL_PATH)
|
|
|
| id2label = model.config.id2label
|
|
|
|
|
| def get_entities_from_sentence(sentence):
|
| words = sentence.split()
|
|
|
| inputs = tokenizer(
|
| words,
|
| return_tensors="pt",
|
| is_split_into_words=True,
|
| truncation=True,
|
| padding=True
|
| )
|
|
|
| with torch.no_grad():
|
| outputs = model(**inputs)
|
|
|
| predictions = torch.argmax(outputs.logits, dim=2)[0].tolist()
|
| word_ids = inputs.word_ids()
|
|
|
| word_labels = []
|
| current_word_id = None
|
| current_labels = []
|
|
|
| for word_id, pred in zip(word_ids, predictions):
|
|
|
| if word_id is None:
|
| continue
|
|
|
| if word_id != current_word_id:
|
| if current_labels:
|
| word_labels.append(id2label[current_labels[0]])
|
|
|
| current_labels = [pred]
|
| current_word_id = word_id
|
| else:
|
| current_labels.append(pred)
|
|
|
| if current_labels:
|
| word_labels.append(id2label[current_labels[0]])
|
|
|
| entities = []
|
|
|
| for i in range(len(words)):
|
| if i < len(word_labels):
|
| if word_labels[i] != "O":
|
| entities.append(words[i])
|
|
|
| return entities |
