import json
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForTokenClassification

sample_inputs = [
    ["Xin chào, tôi tên là Đặng Việt Dũng, tôi sinh ra ở Hòa Bình và đang sinh sống ở Hà Nội. Số điện thoại của tôi là 0948007100 và địa chỉ email của tôi là zzz4everzzz@live.co.uk. Tôi thích chơi game và công nghệ, hiện tôi đang làm việc tại công ty Sabitech, tôi là một lập trình viên. Số căn cước công dân là 113399472. Tôi dễ nổi nóng nhưng có tính cách dễ hòa đồng. Tôi có bằng kỹ sư phần mềm lấy tại trường Đại Học Phương Đông. Là người dân tộc kinh, giới tính nam và không theo tốn giáo nào."],
    ["Tôi muốn mua sản phẩm Nintendo Switch 2, số điện thoại của tôi là 0948007100, vui lòng giao hàng về Dương Nội, Hà Đông, Hà Nội."],
    ["Cho tôi một chiếc áo mầu đen, Vinhome West Point 1, Đỗ Đức Dục, liên hệ 0948007100"]
]

model_name = "mr4/pii-detection-vi"
model = AutoModelForTokenClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)


def ner_predict(text):
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits
    predicted_token_class_ids = torch.argmax(logits, dim=-1)
    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"].squeeze().tolist())
    print("List all token:")
    predicted_entities = [model.config.id2label[id] for id in predicted_token_class_ids.squeeze().tolist()]
    for token, entity in zip(tokens, predicted_entities):
        print(f"  - {token}: {entity}")
    
    final_result = []
    current_entity = []
    current_label = None
    
    for token, entity in zip(tokens, predicted_entities):
        if entity.startswith("B-"):
            if entity[2:] == current_label:
                current_entity.append(token)
            else:
                if current_entity:
                    final_result.append((" ".join(current_entity), current_label))
                current_entity = [token]
                current_label = entity[2:]
        elif entity.startswith("I-") and current_label == entity[2:]:
            current_entity.append(token)
        else:
            if current_entity:
                final_result.append((" ".join(current_entity), current_label))
            current_entity = []
            current_label = None
    if current_entity:
        final_result.append((" ".join(current_entity), current_label))
    return json.dumps(final_result)

# Tạo giao diện
demo = gr.Interface(
    fn=ner_predict,
    inputs=gr.Textbox(lines=5, placeholder="Nhập văn bản tiếng Việt..."),
    outputs=gr.JSON(),
    title="Trích xuất thông tin cá nhân",
    description="Ứng dụng trích xuất thông tin cá nhân trong đoạn văn bản.",
    examples=sample_inputs,
)

# Chạy ứng dụng
demo.launch()