import gradio as gr
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline

# Model details
MODEL_REPO = "xlm-roberta-large-finetuned-conll03-english"

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
model = AutoModelForTokenClassification.from_pretrained(MODEL_REPO)

# Create NER pipeline
ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")

# Define function to extract I-LOC after B-LOC
def extract_locations(text):
    entities = ner_pipeline(text)

    locs = []
    seen_b_loc = False
    for entity in entities:
        if entity["entity_group"] == "LOC":
            if seen_b_loc:
                locs.append(entity["word"])
            seen_b_loc = True
        else:
            seen_b_loc = False

    if locs:
        return ", ".join(locs)
    else:
        return "No I-LOC after B-LOC found."

# Gradio Interface
iface = gr.Interface(
    fn=extract_locations,
    inputs=gr.Textbox(lines=5, placeholder="Enter text here..."),
    outputs="text",
    title="🌍 XLM-RoBERTa Large NER Extractor",
    description="This app finds I-LOC location tags after B-LOC in your input text. Enter a paragraph and see what locations are picked!"
)

# Launch (HF Spaces auto-handles)
iface.launch()