import gradio as gr
import json
from simpletransformers.ner import NERModel
import os
from huggingface_hub import snapshot_download

# Step 1: Download model repo from Hugging Face Hub
repo_path = snapshot_download(repo_id="PixiRus/NER_Model_Version_1")

# Step 2: Define the actual model checkpoint path
model_path = os.path.join(repo_path, "ner_dataset_v1_Model", "checkpoint-119-epoch-1")

# Step 3: Load label mapping from config.json
with open(os.path.join(model_path, "config.json"), "r") as f:
    config = json.load(f)

labels_ = [label for idx, label in sorted(config["id2label"].items(), key=lambda x: int(x[0]))]

# Step 4: Load the NER model
model = NERModel(
    "bert",
    model_path,
    labels=labels_,
    use_cuda=False  # Set to True if running on GPU
)

# Step 5: Define the NER function to return JSON output
def analyze_text(text):
    prediction, _ = model.predict([text])
    tokens = list(prediction[0])
    
    result = []
    for token_dict in tokens:
        for word, label in token_dict.items():
            result.append({
                "word": word,
                "entity": label
            })
    return result

# Step 6: Gradio interface with JSON output
demo = gr.Interface(
    fn=analyze_text,
    inputs=gr.Textbox(lines=5, label="Input Text"),
    outputs=gr.JSON(label="NER Output (JSON)"),
    title="📘 Named Entity Recognition (NER)",
    description="Enter a sentence to extract named entities. The model will return results in JSON format.",
    allow_flagging="never"
)

demo.launch()