ner_app / app.py
vonewman's picture
Create app.py
4d9dd77
raw
history blame
1.13 kB
import streamlit as st
from transformers import AutoModelForTokenClassification, pipeline
# Chargement du modèle de NER
model = AutoModelForTokenClassification.from_pretrained("roberta-base-ner")
# Définition de la pipeline
pipeline = pipeline("ner", model=model)
# Fonction pour traiter un document
def traiter_document(document):
# Tokenisation du document
tokens = model.tokenizer(document)
# Prédiction des entités nommées
predictions = pipeline(tokens)
# Renvoi des entités nommées
return predictions
# Fonction pour charger un document
def charger_document(fichier):
# Ouverture du fichier
with open(fichier, "r") as f:
document = f.read()
return document
# Affichage de l'application
st.title("Named Entity Recognition")
# Chargement du document
document = charger_document(st.file_uploader("Choisissez un document"))
# Traitement du document
predictions = traiter_document(document)
# Affichage des entités nommées
for prediction in predictions:
st.write(f"**Entité:** {prediction['entity_type']}")
st.write(f"**Texte:** {prediction['token']}")