Merwan6
modif
11204e4
raw
history blame
4.94 kB
import torch
import torch.nn.functional as F
from dotenv import load_dotenv
import os
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
#Mapping entre les ID des classes et les labels textuels
id2label = {0: "World", 1: "Sports", 2: "Business", 3: "Sci/Tech"}
def zero_shot_inference(text):
"""
Effectue une classification zero-shot à l'aide du modèle BART MNLI.
Args:
text (str): Texte à classifier.
Returns:
tuple:
- str: Label prédit.
- dict: Dictionnaire {label: score} pour chaque classe.
"""
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
candidate_labels = list(id2label.values())
result = classifier(text, candidate_labels)
prediction = result["labels"][0]
# Formatage des scores avec 4 décimales
scores = {label: float(f"{score:.4f}") for label, score in zip(result["labels"], result["scores"])}
return prediction, scores
def few_shot_inference(text):
"""
Classification few-shot avec FLAN-T5 : génère uniquement le label (World, Sports, etc.).
Args:
text (str): Texte à classifier.
Returns:
tuple:
- str: Label prédit (nettoyé et validé).
- dict: Détails du texte généré brut.
"""
model_name = "google/flan-t5-small"
classifier = pipeline("text2text-generation", model=model_name, max_new_tokens=10)
examples = [
("The president met the UN delegation to discuss global peace.", "World"),
("The football team won their match last night.", "Sports"),
("The company reported a big profit this quarter.", "Business"),
("New research in AI shows promising results.", "Sci/Tech")
]
# Prompt few-shot
prompt = "Classify the following text into one of the following categories: World, Sports, Business, Sci/Tech.\n\n"
for ex_text, ex_label in examples:
prompt += f"Text: {ex_text}\nCategory: {ex_label}\n\n"
prompt += f"Text: {text}\nCategory:"
# Génération
output = classifier(prompt)[0]["generated_text"].strip()
# Nettoyage du label
output_clean = output.split()[0].rstrip(".").capitalize() # ex : "sci/tech." → "Sci/tech"
# Mapping pour être sûr que ça correspond à une catégorie connue
candidate_labels = ["World", "Sports", "Business", "Sci/Tech"]
prediction = next((label for label in candidate_labels if label.lower() in output_clean.lower()), "Unknown")
# Fausse distribution (1.0 pour la classe prédite, 0.0 pour les autres)
scores = {label: 1.0 if label == prediction else 0.0 for label in candidate_labels}
return prediction, scores
def base_model_inference(text):
"""
Utilise un modèle BERT préentraîné sur AG News (sans fine-tuning personnalisé).
Args:
text (str): Texte à classifier.
Returns:
tuple:
- str: Label prédit.
- dict: Scores softmax par classe.
"""
model_name = "textattack/bert-base-uncased-ag-news"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
#Encodage du texte
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
#Prédiction sans calcul de gradients
with torch.no_grad():
outputs = model(**inputs)
#Calcul des probabilités avec softmax
probs = F.softmax(outputs.logits, dim=-1)[0].cpu().numpy()
pred_id = probs.argmax()
prediction = id2label[pred_id]
scores = {id2label[i]: float(f"{p:.4f}") for i, p in enumerate(probs)}
return prediction, scores
def fine_tuned_inference(text, model_path="Merwan611/agnews-finetuned-bert"):
"""
Utilise un modèle BERT fine-tuné personnalisé sur AG News, avec authentification Hugging Face si nécessaire.
Args:
text (str): Texte à classifier.
model_path (str): Nom du modèle Hugging Face ou chemin local.
Returns:
tuple:
- str: Label prédit.
- dict: Scores softmax par classe.
"""
#Récupération du token d'auth depuis les variables d'environnement
token = os.getenv("CLE")
tokenizer = AutoTokenizer.from_pretrained(model_path, use_auth_token=token)
model = AutoModelForSequenceClassification.from_pretrained(model_path)
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
with torch.no_grad():
outputs = model(**inputs)
probs = F.softmax(outputs.logits, dim=-1)[0].cpu().numpy()
pred_id = probs.argmax()
prediction = id2label[pred_id]
scores = {id2label[i]: float(f"{p:.4f}") for i, p in enumerate(probs)}
return prediction, scores