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']}")