from transformers import AutoTokenizer, AutoModelForTokenClassification
import torch
import streamlit as st

model = AutoModelForTokenClassification.from_pretrained('./ner_model')
tokenizer = AutoTokenizer.from_pretrained('./ner_model')
label_list = model.config.id2label


def predict_entities(text):
    tokens = tokenizer(text, return_tensors='pt', truncation=True, is_split_into_words=False)
    with torch.no_grad():
        outputs = model(**tokens)

    predictions = torch.argmax(outputs.logits, dim=2)
    tokens = tokens['input_ids'][0]
    prediction_ids = predictions[0]

    result = []
    for token_id, pred_id in zip(tokens, prediction_ids):
        token = tokenizer.decode([token_id])
        label = label_list[pred_id.item()]
        if token not in ['[CLS]', '[SEP]', '[PAD]']:
            result.append((token, label))
    return result

st.set_page_config(page_title="NER App", layout="wide")

st.title("🧠 Named Entity Recognition (NER) with BERT")
text_input = st.text_area("Enter text to analyze:")

if st.button("Analyze"):
    if text_input:
        results = predict_entities(text_input)
        st.markdown("### 🔍 Extracted Entities:")
        for word, label in results:
            if label != "O":
                st.markdown(f"**{word}** → `{label}`")
    else:
        st.warning("Please enter some text.")