# -*- coding: utf-8 -*- import spacy import torch from spacy import displacy import spacy_curated_transformers from propp_fr import ( generate_tokens_df, load_mentions_detection_model, load_coreference_resolution_model, load_tokenizer_and_embedding_model, get_embedding_tensor_from_tokens_df, generate_entities_df, add_features_to_entities, perform_coreference ) # ----------------------------- # MODELS # ----------------------------- NER_MODEL = "AntoineBourgois/propp-fr_NER_camembert-large_FAC_GPE_LOC_PER_TIME_VEH" COREF_MODEL = "AntoineBourgois/propp-fr_coreference-resolution_camembert-large_PER" PALETTE = [ "#f4a261", "#2a9d8f", "#e76f51", "#457b9d", "#8ecae6", "#e9c46a", "#06d6a0", "#ef476f", "#a8dadc", "#264653" ] # ----------------------------- # CORÉFÉRENCE RULE-BASED (non-PER) # ----------------------------- def rule_based_coref_for_non_per(entities_df, coref_id_start): """ Assigne des COREF IDs aux entités non-PER par correspondance de texte (par type). Les entités de même texte (normalisé) et de même type sont regroupées dans le même cluster. """ non_per_cats = [c for c in entities_df["cat"].dropna().unique() if c != "PER"] next_id = coref_id_start for cat in non_per_cats: cat_mask = entities_df["cat"] == cat cat_entities = entities_df[cat_mask] text_to_ids = {} for idx, row in cat_entities.iterrows(): key = str(row.get("text", "")).lower().strip() if not key: key = f"__unknown_{idx}__" text_to_ids.setdefault(key, []).append(idx) for ids in text_to_ids.values(): entities_df.loc[ids, "COREF"] = next_id next_id += 1 return entities_df # ----------------------------- # FASTCOREF (LingMess) # ----------------------------- def load_lingmess_model(device="cuda:0"): from fastcoref import LingMessCoref return LingMessCoref(device=device) def merge_fastcoref(entities_df, tokens_df, text, lingmess_model): """ Enrichit les COREF via LingMess (coréférences inter-types et non-PER). Clusters purement PER : propp-fr garde la priorité. Clusters non-PER ou mixtes : LingMess fusionne les IDs. Entités partageant le même COREF ID (même cross-type) → même couleur. """ preds = lingmess_model.predict(texts=[text]) lingmess_clusters = preds[0].get_clusters(as_strings=False) b2c = build_byte_to_char_map(text) entity_indices = list(entities_df.index) entity_char_spans = [] for _, row in entities_df.iterrows(): st = int(row["start_token"]) et = int(row["end_token"]) b_start = int(tokens_df.loc[ tokens_df["token_ID_within_document"] == st, "byte_onset" ].values[0]) b_end = int(tokens_df.loc[ tokens_df["token_ID_within_document"] == et, "byte_offset" ].values[0]) entity_char_spans.append((b2c.get(b_start, b_start), b2c.get(b_end, b_end))) def find_entity_for_span(c_start, c_end): best_idx, best_overlap = None, 0 for i, (es, ee) in enumerate(entity_char_spans): overlap = min(c_end, ee) - max(c_start, es) if overlap > best_overlap: best_overlap = overlap best_idx = i return best_idx if best_overlap > 0 else None next_id = int(entities_df["COREF"].dropna().max()) + 1 if entities_df["COREF"].notna().any() else 0 for cluster in lingmess_clusters: matched = [] for c_start, c_end in cluster: local_idx = find_entity_for_span(c_start, c_end) if local_idx is not None: df_idx = entity_indices[local_idx] if df_idx not in matched: matched.append(df_idx) if len(matched) < 2: continue cats = entities_df.loc[matched, "cat"].tolist() unique_cats = set(cats) # propp-fr gère mieux les clusters PER purs if unique_cats == {"PER"}: continue # LingMess (entraîné sur l'anglais) génère trop de faux positifs # cross-type sur le français → on n'accepte que les fusions intra-type non-PER if len(unique_cats) > 1: continue existing = entities_df.loc[matched, "COREF"].dropna() target_id = int(existing.iloc[0]) if len(existing) > 0 else next_id if len(existing) == 0: next_id += 1 entities_df.loc[matched, "COREF"] = target_id return entities_df # ----------------------------- # CHARGEMENT DES MODÈLES (une seule fois) # ----------------------------- def load_all_models(with_lingmess=True, device="cuda:0"): """ Charge tous les modèles une seule fois. Réutilisable pour un serveur : évite de recharger CamemBERT-large à chaque requête. """ spacy.prefer_gpu() nlp = spacy.load("fr_dep_news_trf") mentions_model = load_mentions_detection_model(NER_MODEL, force_download=False) base = mentions_model["base_model_name"] tokenizer, embedding_model = load_tokenizer_and_embedding_model(base) coref_model = load_coreference_resolution_model(COREF_MODEL, force_download=False) lingmess_model = None if with_lingmess: lingmess_model = load_lingmess_model(device=device) return { "nlp": nlp, "mentions_model": mentions_model, "base": base, "tokenizer": tokenizer, "embedding_model": embedding_model, "coref_model": coref_model, "lingmess_model": lingmess_model, } # ----------------------------- # PIPELINE PROPP-FR # ----------------------------- def run_pipeline(text, lingmess_model=None, models=None): print("GPU:", torch.cuda.is_available()) # Réutilise les modèles préchargés si fournis, sinon charge à la volée if models is None: models = load_all_models(with_lingmess=False) if lingmess_model is None: lingmess_model = models.get("lingmess_model") nlp = models["nlp"] mentions_model = models["mentions_model"] base = models["base"] tokenizer = models["tokenizer"] embedding_model = models["embedding_model"] coref_model = models["coref_model"] # TOKENS tokens_df = generate_tokens_df(text, nlp) # NER emb = get_embedding_tensor_from_tokens_df( text, tokens_df, tokenizer, embedding_model, mini_batch_size=32, subword_pooling_strategy=mentions_model["subword_pooling_strategy"] ) entities_df = generate_entities_df(tokens_df, emb, mentions_model, batch_size=32) entities_df = add_features_to_entities(entities_df, tokens_df) # COREF ML (PER uniquement — seul modèle disponible) if coref_model["base_model_name"] != base: tokenizer, embedding_model = load_tokenizer_and_embedding_model( coref_model["base_model_name"] ) emb = get_embedding_tensor_from_tokens_df(text, tokens_df, tokenizer, embedding_model) entities_df = perform_coreference( entities_df=entities_df, tokens_embedding_tensor=emb, coreference_resolution_model=coref_model, batch_size=50000, propagate_coref=True, rule_based_postprocess=False ) # COREF rule-based pour FAC, GPE, LOC, TIME, VEH (groupement par texte identique) next_id = int(entities_df["COREF"].dropna().max()) + 1 if entities_df["COREF"].notna().any() else 0 entities_df = rule_based_coref_for_non_per(entities_df, next_id) # COREF LingMess : coréférences inter-types et non-PER (optionnel) if lingmess_model is not None: print("Fusion LingMess...") entities_df = merge_fastcoref(entities_df, tokens_df, text, lingmess_model) return tokens_df, entities_df # ----------------------------- # CONVERSION BYTE → CHAR # ----------------------------- def build_byte_to_char_map(text): table = {} b = 0 for i, ch in enumerate(text): table[b] = i b += len(ch.encode("utf-8")) table[b] = len(text) return table # ----------------------------- # DISPLACY BUILDER # ----------------------------- def build_displacy_entities(text, tokens_df, entities_df): ents = [] for _, row in entities_df.iterrows(): if row["COREF"] is None: continue start_token = int(row["start_token"]) end_token = int(row["end_token"]) c_start = int(tokens_df.loc[ tokens_df["token_ID_within_document"] == start_token, "byte_onset" ].values[0]) c_end = int(tokens_df.loc[ tokens_df["token_ID_within_document"] == end_token, "byte_offset" ].values[0]) cat = str(row.get("cat", "ENT")) ents.append({ "start": c_start, "end": c_end, "label": f"{cat}_{int(row['COREF'])}", }) ents.sort(key=lambda e: e["start"]) return ents # ----------------------------- # COLORS # ----------------------------- def build_colors(entities_df): # Une couleur par COREF ID — partagée entre tous les types du même cluster cores = sorted(entities_df["COREF"].dropna().astype(int).unique()) coref_to_color = {cid: PALETTE[i % len(PALETTE)] for i, cid in enumerate(cores)} valid = entities_df[entities_df["COREF"].notna()][["COREF", "cat"]].drop_duplicates() return { f"{row['cat']}_{int(row['COREF'])}": coref_to_color[int(row['COREF'])] for _, row in valid.iterrows() } # ----------------------------- # VISUALISATION HTML # ----------------------------- def export_html(text, tokens_df, entities_df, output="coref.html"): ents = build_displacy_entities(text, tokens_df, entities_df) colors = build_colors(entities_df) html = displacy.render( { "text": text, "ents": ents, "title": "Propp-fr Coreference (corrected)" }, style="ent", manual=True, options={"colors": colors}, page=True ) with open(output, "w", encoding="utf-8") as f: f.write(html) print("✔ HTML généré :", output) # ----------------------------- # DEMO # ----------------------------- if __name__ == "__main__": texte = """Léo est dans la cuisine. Dans l'après-midi, il a des invités. Il décide de faire une salade composée pour ses invités. Sur une table, il prépare les ingrédients et le matériel. Il égoutte le maïs avec la passoire. Il lave, il épluche et il râpe des carottes. Il coupe le gruyère en petits cubes. Il retire la peau et le noyau de l'avocat. Il coupe l'avocat en fines lamelles. Il fait ensuite une vinaigrette avec l'huile, le vinaigre, le sel et le poivre. Il verse la vinaigrette dans un saladier avec les carottes, le maïs, l'avocat et le gruyère. Il mélange. La salade composée est prête.""" lingmess = load_lingmess_model(device="cuda:0") tokens_df, entities_df = run_pipeline(texte, lingmess_model=lingmess) export_html(texte, tokens_df, entities_df) print(entities_df[["text", "cat", "prop", "COREF"]].to_string())