Spaces:
Running
Running
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() | |