import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification
from seqeval.metrics.sequence_labeling import get_entities

tokenizer = AutoTokenizer.from_pretrained("./model", use_fast = False)
model = AutoModelForTokenClassification.from_pretrained("./model")

def get_words(tokens:list[str]):
    return tokenizer.convert_tokens_to_string(tokens)

def format_result(tokens:list[str], labels:list[str]):
    """Định dạng kết quả NER thành HTML với các thực thể được đánh dấu."""
    
    formatted_output = ""
    current_position = 0
    start = 0
    end = 0
    EntitySpan = '<span class="entity" data-entity="{label}">{word} <span class="entity" label-entity="{label}">{label}</span></span>'
    
    for label, start, end in get_entities(labels):
        
        end += 1
        
        entity = {
            'word':get_words(tokens[start:end]),
            'label':label
        }
        
        # Thêm phần văn bản trước thực thể (nếu có)
        if start > current_position:
            formatted_output += get_words(tokens[current_position: start])

        # Thêm thực thể với thẻ span và nhãn
        formatted_output += EntitySpan.format(**entity)

        # Cập nhật vị trí hiện tại
        current_position = end

    # Thêm phần văn bản còn lại sau thực thể cuối cùng (nếu có)
    if current_position < len(tokens):
        formatted_output += get_words(tokens[current_position:])

    return formatted_output

def inference(text:str):
    
    inputs = tokenizer(text, return_tensors="pt",add_special_tokens=False)

    with torch.no_grad():
        logits = model(**inputs).logits
        predictions = torch.argmax(logits, dim=2)
        labels = [model.config.id2label[t.item()] for t in predictions[0]]
        
    tokens = [t.detach().numpy() for t in inputs['input_ids']]
    tokens = tokenizer.convert_ids_to_tokens(tokens[0])
    return format_result(tokens, labels)