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

# Library for Downloading custom model from HuggingFace Hub
from huggingface_hub import snapshot_download

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

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

# Example Text

example_sent = (
    '''LE BLOND (Guillaume)  - L’Artillerie raisonnée contenant la description et l’usage des différentes bouches à feu... La Théorie & la pratique des mines, & du jet des bombes... / par M. Le Blond, ... - À Paris, chez CharL. Ant. Jombert, 1761. - XXII-579-[4] p.-[30] f. de dépl. ; in-8 (20 cm)  Rel. veau marbré  Sig. à8, b4, A-Z8, Aa-Nn8, Oo4  Rx 216  Artillerie'''
)

# 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
)

# Function to process and highlight NER predictions
def analyze_text(text):
    prediction, _ = model.predict([text])
    tokens = list(prediction[0])
    
    highlighted = []
    for token_dict in tokens:
        for word, label in token_dict.items():
            tag = label if label != "O" else None
            highlighted.append((word + " ", tag))
    
    return highlighted

# Build the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## 🤖 AI Based NER Model")

    input_text = gr.Textbox(lines=4, label="Enter text", value=example_sent)
    analyze_btn = gr.Button("Run NER")
    output = gr.HighlightedText(label="NER Output")

    analyze_btn.click(analyze_text, inputs=input_text, outputs=output)

demo.launch()