import gradio as gr
import pdfplumber
from docx import Document
from transformers import pipeline
import os

# Charger le modèle NER
ner_pipeline = pipeline("ner", model="yashpwr/resume-ner-bert-v2")

def extract_text_from_pdf(file):
    with pdfplumber.open(file) as pdf:
        text = " ".join([page.extract_text() or "" for page in pdf.pages])
    return text

def extract_text_from_docx(file):
    doc = Document(file)
    text = " ".join([para.text for para in doc.paragraphs])
    return text

def extract_text(file):
    file_ext = os.path.splitext(file.name)[1].lower()
    if file_ext == ".pdf":
        return extract_text_from_pdf(file.name)
    elif file_ext == ".docx":
        return extract_text_from_docx(file.name)
    elif file_ext == ".txt":
        with open(file.name, "r") as f:
            return f.read()
    else:
        return "Format non supporté. Utilisez PDF, DOCX ou TXT."

def parse_resume(file):
    if file is None:
        return "Veuillez uploader un CV"
    text = extract_text(file)
    if not text.strip():
        return "Aucun texte trouvé dans le fichier"
    
    # Limiter la taille (le modèle a une limite de 512 tokens)
    text = text[:2000]
    entities = ner_pipeline(text)
    
    # Formater les résultats
    results = {}
    for ent in entities:
        entity_type = ent['entity'].replace("B-", "").replace("I-", "")
        word = ent['word']
        if entity_type not in results:
            results[entity_type] = []
        if word not in results[entity_type]:  # éviter doublons
            results[entity_type].append(word)
    
    # Convertir en texte lisible
    output = "=== ENTITÉS EXTRAITES ===\n"
    for k, v in results.items():
        output += f"\n{k}: {', '.join(v)}\n"
    return output

# Interface Gradio
iface = gr.Interface(
    fn=parse_resume,
    inputs=gr.File(label="Téléchargez votre CV (PDF, DOCX, TXT)"),
    outputs=gr.Textbox(label="Résultat", lines=20),
    title="Parseur de CV avec NER",
    description="Extrait automatiquement le nom, email, téléphone, compétences, diplômes, etc."
)

iface.launch()