may24c/insertion-paragraphe

1import torch2from transformers import AutoTokenizer, AutoModelForSequenceClassification3 4# ========================5# CONFIGURATION DU MODÈLE ROBERTA6# ========================7 8ROBERTA_PATH = "may24c/model_roberta"9MAX_LEN_ROBERTA = 25610 11# Tokenizer et modèle12tokenizer_roberta = AutoTokenizer.from_pretrained(ROBERTA_PATH)13model_roberta = AutoModelForSequenceClassification.from_pretrained(ROBERTA_PATH)14 15# Device (GPU si dispo sinon CPU)16device = torch.device("cuda" if torch.cuda.is_available() else "cpu")17model_roberta.to(device)18model_roberta.eval()19 20# ========================21# FONCTION PREDICTION22# ========================23def predict_insertion(text, paragraph):24    """25    Retourne la meilleure position pour insérer un paragraphe selon nôtre model RoBERTa26    """27    paragraphs = text.split("\n")28    best_idx = 029    best_score = -float("inf")30 31    for i in range(len(paragraphs) + 1):32        prev_text = "\n".join(paragraphs[:i])33        next_text = "\n".join(paragraphs[i:])34 35        # Préparation input RoBERTa36        inputs = tokenizer_roberta(37            text=str(paragraph),38            text_pair=str(prev_text) + f" {tokenizer_roberta.sep_token} " + str(next_text),39            return_tensors="pt",40            truncation="longest_first",41            max_length=MAX_LEN_ROBERTA,42        )43 44        inputs = {k: v.to(device) for k, v in inputs.items()}45 46        with torch.no_grad():47            outputs = model_roberta(**inputs)48 49        probs = torch.softmax(outputs.logits, dim=1)50        score = probs[0][1].item()51 52        if score > best_score:53            best_score = score54            best_idx = i55 56    paragraphs.insert(best_idx, paragraph)57    new_text = "\n".join(paragraphs)58 59    return {60        "position": best_idx,61        "score": best_score,62        "text_modified": new_text63    }