import os
from transformers import AutoTokenizer, AutoModelForTokenClassification
import torch

MODEL_PATH = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "..", "models", "ner_model")
)

# Fallback to the Hugging Face hub repository if the local model files are not present
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