nlp-project / app.py
kohlin's picture
Initial Commit
978a5b4
raw
history blame
617 Bytes
import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification
# Load fine-tuned model
model_path = "./ner_model"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForTokenClassification.from_pretrained(model_path)
# Create NER pipeline
ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)
def ner_prediction(text):
entities = ner_pipeline(text)
return {e["word"]: e["entity"] for e in entities}
# Gradio UI
iface = gr.Interface(fn=ner_prediction, inputs="text", outputs="label")
iface.launch(server_name="0.0.0.0", server_port=7860)