scwaterbear's picture
Update app.py
3d49345 verified
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()