import gradio as gr
from transformers import AutoTokenizer, AutoModelForTokenClassification
from transformers import pipeline

# Charger le tokenizer et le modèle pour la reconnaissance d'entités nommées (NER)
tokenizer = AutoTokenizer.from_pretrained("sp01/ner_automotive")
model = AutoModelForTokenClassification.from_pretrained("sp01/ner_automotive")

# Créer un pipeline pour la NER
ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)

# Fonction pour traiter l'entrée
def extract_entities(text):
    entities = ner_pipeline(text)
    return "\n".join([f"Entity: {entity['word']}, Label: {entity['entity']}, Score: {entity['score']}" for entity in entities])

# Interface Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Extraction d'entités automobiles")
    
    # Composants de l'interface
    text_input = gr.Textbox(label="Entrez votre texte", placeholder="Ex: Tesla Model S est une voiture électrique.")
    output_text = gr.Textbox(label="Entités extraites", lines=5)
    
    # Lier le bouton à la fonction
    text_input.submit(extract_entities, inputs=text_input, outputs=output_text)

demo.launch()