Spaces:

BHM03
/

amane-api

Running

amane-api / predictor.py

Initial AMANE API deploy — FastAPI + Docker

d40eb53 6 months ago

19.1 kB

import os
import json
import joblib
import numpy as np
from utils import clean_text, LABEL_COLS

# torch et transformers sont optionnels — uniquement requis pour le Teacher XLM-R
try:
    import torch
    import torch.nn as nn
    from transformers import AutoTokenizer, AutoModel
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False


# ── HuggingFace Hub — fallback si fichiers absents en local ─────────────────
HF_REPO_ID = 'BHM03/amane-models'

def _resolve(local_path: str, hf_filename: str) -> str:
    """Retourne le chemin local si le fichier existe, sinon télécharge depuis HF Hub."""
    if os.path.exists(local_path):
        return local_path
    try:
        from huggingface_hub import hf_hub_download
        return hf_hub_download(
            repo_id=HF_REPO_ID,
            filename=hf_filename,
            token=os.environ.get('HF_TOKEN'),
        )
    except Exception as e:
        raise FileNotFoundError(
            f'Fichier introuvable localement ({local_path}) '
            f'et téléchargement HF Hub échoué : {e}'
        )


# ── Chemins modèles ──────────────────────────────────────────────────────────
BENCHMARK_MODEL_PATH    = 'models/ovr_balanced.joblib'
BENCHMARK_TFIDF_PATH    = 'models/tfidf_vectorizer.joblib'
TEACHER_CKPT_PATH       = 'models/teacher/teacher_best.pt'
TEACHER_TOKENIZER_PATH  = 'models/teacher/tokenizer'
TEACHER_THRESH_PATH     = 'models/teacher/thresholds.json'
TEACHER_STAR_CKPT_PATH  = 'models/teacher_star/teacher_star.pt'
TEACHER_STAR_THRESH_PATH= 'models/teacher_star/threshold_darija.json'
STUDENT_CKPT_PATH       = 'models/student/student_best.pt'
STUDENT_THRESH_PATH     = 'models/student/threshold_student.json'

# ── Hyperparamètres d'inférence Teacher (cohérents avec l'entraînement) ─────
TEACHER_MAX_LEN      = 256
TEACHER_STAR_MAX_LEN = 128
STUDENT_MAX_LEN      = 128
TEACHER_MODEL_HF     = 'xlm-roberta-base'
STUDENT_MODEL_HF     = 'distilbert-base-multilingual-cased'


# ════════════════════════════════════════════════════════════════════════════
# BENCHMARK — OneVsRest + LogisticRegression
# ════════════════════════════════════════════════════════════════════════════
class BenchmarkPredictor:
    def __init__(self):
        self.model    = None
        self.tfidf    = None
        self.loaded   = False
        self.error    = None

    def load(self):
        try:
            self.model  = joblib.load(_resolve(BENCHMARK_MODEL_PATH, 'ovr_balanced.joblib'))
            self.tfidf  = joblib.load(_resolve(BENCHMARK_TFIDF_PATH, 'tfidf_vectorizer.joblib'))
            self.loaded = True
        except Exception as e:
            self.error  = str(e)
            self.loaded = False

    def predict(self, text: str) -> dict:
        if not self.loaded:
            raise RuntimeError(f'Benchmark non chargé : {self.error}')
        cleaned  = clean_text(text).lower()
        vec      = self.tfidf.transform([cleaned])
        probs    = self.model.predict_proba(vec)[0]
        return {col: float(round(p, 4)) for col, p in zip(LABEL_COLS, probs)}

    def predict_batch(self, texts: list) -> list[dict]:
        if not self.loaded:
            raise RuntimeError(f'Benchmark non chargé : {self.error}')
        cleaned = [clean_text(t).lower() for t in texts]
        vecs    = self.tfidf.transform(cleaned)
        probs   = self.model.predict_proba(vecs)
        return [
            {col: float(round(p, 4)) for col, p in zip(LABEL_COLS, row)}
            for row in probs
        ]


# ════════════════════════════════════════════════════════════════════════════
# XLM-R TEACHER
# ════════════════════════════════════════════════════════════════════════════
class TeacherPredictor:
    def __init__(self):
        self.model      = None
        self.tokenizer  = None
        self.thresholds = None
        self.loaded     = False
        self.error      = None
        self.device     = None

    def load(self):
        if not TORCH_AVAILABLE:
            self.error  = (
                'PyTorch / Transformers non installés. '
                'Lance : pip install torch transformers sentencepiece'
            )
            self.loaded = False
            return

        try:
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

            # 2. Charger le tokenizer — on utilise directement 'xlm-roberta-base'
            # car le tokenizer ne change pas lors du fine-tuning, et les fichiers
            # locaux sont incomplets (sentencepiece.bpe.model manquant).
            self.tokenizer = AutoTokenizer.from_pretrained(TEACHER_MODEL_HF)

            # 3. Reconstruire l'architecture (identique au notebook d'entraînement)
            class TeacherModel(nn.Module):
                def __init__(self, model_name: str, num_labels: int = 6, dropout: float = 0.1):
                    super().__init__()
                    self.encoder    = AutoModel.from_pretrained(model_name)
                    hidden_size     = self.encoder.config.hidden_size
                    self.dropout    = nn.Dropout(dropout)
                    self.classifier = nn.Linear(hidden_size, num_labels)

                def mean_pool(self, token_embeds, attention_mask):
                    mask   = attention_mask.unsqueeze(-1).expand(token_embeds.size()).float()
                    summed = (token_embeds * mask).sum(dim=1)
                    counts = mask.sum(dim=1).clamp(min=1e-9)
                    return summed / counts

                def forward(self, input_ids, attention_mask):
                    outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
                    pooled  = self.mean_pool(outputs.last_hidden_state, attention_mask)
                    pooled  = self.dropout(pooled)
                    return self.classifier(pooled)

            # 4. Charger le checkpoint
            ckpt  = torch.load(_resolve(TEACHER_CKPT_PATH, 'teacher_best.pt'), map_location=self.device, weights_only=False)
            model = TeacherModel(TEACHER_MODEL_HF, num_labels=len(LABEL_COLS))
            model.load_state_dict(ckpt['model_state'])
            model.to(self.device)
            model.eval()
            self.model = model

            # 5. Charger les seuils tunés
            if os.path.exists(TEACHER_THRESH_PATH):
                with open(TEACHER_THRESH_PATH) as f:
                    self.thresholds = json.load(f)
            else:
                # Fallback : seuil 0.5 par défaut
                self.thresholds = {col: 0.5 for col in LABEL_COLS}

            self.loaded = True

        except Exception as e:
            self.error  = f'{type(e).__name__}: {e}'
            self.loaded = False

    # ── Inférence : probabilités ───────────────────────────────────────────
    def predict(self, text: str) -> dict:
        """Retourne les probabilités sigmoïdes par label."""
        if not self.loaded:
            raise RuntimeError(f'Teacher non chargé : {self.error}')

        cleaned  = clean_text(text)
        encoding = self.tokenizer(
            cleaned,
            max_length=TEACHER_MAX_LEN,
            padding='max_length',
            truncation=True,
            return_tensors='pt',
        )
        input_ids = encoding['input_ids'].to(self.device)
        attn_mask = encoding['attention_mask'].to(self.device)

        with torch.no_grad():
            logits = self.model(input_ids, attn_mask)
            probs  = torch.sigmoid(logits).cpu().numpy()[0]

        return {col: float(round(p, 4)) for col, p in zip(LABEL_COLS, probs)}

    # ── Inférence : labels binaires (avec seuils tunés) ────────────────────
    def predict_with_labels(self, text: str) -> dict:
        """
        Retourne probabilités + labels binaires selon les seuils tunés.
        Format : {'probs': {...}, 'labels': {...}, 'thresholds': {...}}
        """
        probs  = self.predict(text)
        labels = {col: int(probs[col] >= self.thresholds[col]) for col in LABEL_COLS}
        return {
            'probs'      : probs,
            'labels'     : labels,
            'thresholds' : self.thresholds,
        }

    # ── Inférence : batch ──────────────────────────────────────────────────
    def predict_batch(self, texts: list) -> list[dict]:
        """Inférence batch — plus efficace que predict() en boucle."""
        if not self.loaded:
            raise RuntimeError(f'Teacher non chargé : {self.error}')

        cleaned = [clean_text(t) for t in texts]
        encoding = self.tokenizer(
            cleaned,
            max_length=TEACHER_MAX_LEN,
            padding='max_length',
            truncation=True,
            return_tensors='pt',
        )
        input_ids = encoding['input_ids'].to(self.device)
        attn_mask = encoding['attention_mask'].to(self.device)

        with torch.no_grad():
            logits = self.model(input_ids, attn_mask)
            probs  = torch.sigmoid(logits).cpu().numpy()

        return [
            {col: float(round(p, 4)) for col, p in zip(LABEL_COLS, row)}
            for row in probs
        ]


# ════════════════════════════════════════════════════════════════════════════
# XLM-R TEACHER* — adapté Darija (fine-tuné sur DarLoad, loss masquée)
# ════════════════════════════════════════════════════════════════════════════
class TeacherStarPredictor:
    def __init__(self):
        self.model          = None
        self.tokenizer      = None
        self.threshold      = 0.45   # seuil tunée sur DarLoad
        self.loaded         = False
        self.error          = None
        self.device         = None

    def load(self):
        if not TORCH_AVAILABLE:
            self.error  = 'PyTorch / Transformers non installés.'
            self.loaded = False
            return

        try:
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
            self.tokenizer = AutoTokenizer.from_pretrained(TEACHER_MODEL_HF)

            class TeacherModel(nn.Module):
                def __init__(self, model_name, num_labels=6, dropout=0.1):
                    super().__init__()
                    self.encoder    = AutoModel.from_pretrained(model_name)
                    self.dropout    = nn.Dropout(dropout)
                    self.classifier = nn.Linear(self.encoder.config.hidden_size, num_labels)

                def mean_pool(self, token_embeds, attention_mask):
                    mask   = attention_mask.unsqueeze(-1).expand(token_embeds.size()).float()
                    summed = (token_embeds * mask).sum(dim=1)
                    return summed / mask.sum(dim=1).clamp(min=1e-9)

                def forward(self, input_ids, attention_mask):
                    out    = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
                    pooled = self.mean_pool(out.last_hidden_state, attention_mask)
                    return self.classifier(self.dropout(pooled))

            ckpt  = torch.load(_resolve(TEACHER_STAR_CKPT_PATH, 'teacher_star.pt'), map_location=self.device, weights_only=False)
            model = TeacherModel(TEACHER_MODEL_HF, num_labels=len(LABEL_COLS))
            model.load_state_dict(ckpt['model_state'])
            model.to(self.device)
            model.eval()
            self.model = model

            if os.path.exists(TEACHER_STAR_THRESH_PATH):
                with open(TEACHER_STAR_THRESH_PATH) as f:
                    self.threshold = json.load(f).get('toxic_darija', 0.45)

            self.loaded = True

        except Exception as e:
            self.error  = f'{type(e).__name__}: {e}'
            self.loaded = False

    def predict(self, text: str) -> dict:
        if not self.loaded:
            raise RuntimeError(f'Teacher* non chargé : {self.error}')
        cleaned  = clean_text(text)
        encoding = self.tokenizer(
            cleaned, max_length=TEACHER_STAR_MAX_LEN,
            padding='max_length', truncation=True, return_tensors='pt',
        )
        with torch.no_grad():
            logits = self.model(
                encoding['input_ids'].to(self.device),
                encoding['attention_mask'].to(self.device),
            )
            probs = torch.sigmoid(logits).cpu().numpy()[0]
        return {col: float(round(p, 4)) for col, p in zip(LABEL_COLS, probs)}

    def predict_batch(self, texts: list) -> list[dict]:
        if not self.loaded:
            raise RuntimeError(f'Teacher* non chargé : {self.error}')
        cleaned  = [clean_text(t) for t in texts]
        encoding = self.tokenizer(
            cleaned, max_length=TEACHER_STAR_MAX_LEN,
            padding='max_length', truncation=True, return_tensors='pt',
        )
        with torch.no_grad():
            logits = self.model(
                encoding['input_ids'].to(self.device),
                encoding['attention_mask'].to(self.device),
            )
            probs = torch.sigmoid(logits).cpu().numpy()
        return [{col: float(round(p, 4)) for col, p in zip(LABEL_COLS, row)} for row in probs]


# ════════════════════════════════════════════════════════════════════════════
# FACTORY — chargement unique via st.cache_resource
# ════════════════════════════════════════════════════════════════════════════
def load_benchmark() -> BenchmarkPredictor:
    p = BenchmarkPredictor()
    p.load()
    return p

def load_teacher() -> TeacherPredictor:
    p = TeacherPredictor()
    p.load()
    return p

def load_teacher_star() -> TeacherStarPredictor:
    p = TeacherStarPredictor()
    p.load()
    return p


# ════════════════════════════════════════════════════════════════════════════
# STUDENT — DistilBERT multilingue spécialisé Darija
# ════════════════════════════════════════════════════════════════════════════
class StudentPredictor:
    def __init__(self):
        self.model     = None
        self.tokenizer = None
        self.threshold = 0.6
        self.loaded    = False
        self.error     = None
        self.device    = None

    def load(self):
        if not TORCH_AVAILABLE:
            self.error  = 'PyTorch / Transformers non installés.'
            self.loaded = False
            return

        try:
            self.device    = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
            self.tokenizer = AutoTokenizer.from_pretrained(STUDENT_MODEL_HF)

            class StudentModel(nn.Module):
                def __init__(self, model_name, num_labels=6, dropout=0.1):
                    super().__init__()
                    self.encoder    = AutoModel.from_pretrained(model_name)
                    self.dropout    = nn.Dropout(dropout)
                    self.classifier = nn.Linear(self.encoder.config.hidden_size, num_labels)

                def forward(self, input_ids, attention_mask):
                    out    = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
                    pooled = out.last_hidden_state[:, 0, :]  # token [CLS]
                    return self.classifier(self.dropout(pooled))

            ckpt  = torch.load(_resolve(STUDENT_CKPT_PATH, 'student_best.pt'), map_location=self.device, weights_only=False)
            model = StudentModel(STUDENT_MODEL_HF, num_labels=len(LABEL_COLS))
            model.load_state_dict(ckpt['model_state'])
            model.to(self.device)
            model.eval()
            self.model = model

            if os.path.exists(STUDENT_THRESH_PATH):
                with open(STUDENT_THRESH_PATH) as f:
                    self.threshold = json.load(f).get('toxic_darija', 0.35)

            self.loaded = True

        except Exception as e:
            self.error  = f'{type(e).__name__}: {e}'
            self.loaded = False

    def predict(self, text: str) -> dict:
        if not self.loaded:
            raise RuntimeError(f'Student non chargé : {self.error}')
        cleaned  = clean_text(text)
        encoding = self.tokenizer(
            cleaned, max_length=STUDENT_MAX_LEN,
            padding='max_length', truncation=True, return_tensors='pt',
        )
        with torch.no_grad():
            logits = self.model(
                encoding['input_ids'].to(self.device),
                encoding['attention_mask'].to(self.device),
            )
            probs = torch.sigmoid(logits).cpu().numpy()[0]
        return {col: float(round(p, 4)) for col, p in zip(LABEL_COLS, probs)}

    def predict_batch(self, texts: list) -> list[dict]:
        if not self.loaded:
            raise RuntimeError(f'Student non chargé : {self.error}')
        cleaned  = [clean_text(t) for t in texts]
        encoding = self.tokenizer(
            cleaned, max_length=STUDENT_MAX_LEN,
            padding='max_length', truncation=True, return_tensors='pt',
        )
        with torch.no_grad():
            logits = self.model(
                encoding['input_ids'].to(self.device),
                encoding['attention_mask'].to(self.device),
            )
            probs = torch.sigmoid(logits).cpu().numpy()
        return [{col: float(round(p, 4)) for col, p in zip(LABEL_COLS, row)} for row in probs]


def load_student() -> StudentPredictor:
    p = StudentPredictor()
    p.load()
    return p

	import os
	import json
	import joblib
	import numpy as np
	from utils import clean_text, LABEL_COLS

	# torch et transformers sont optionnels — uniquement requis pour le Teacher XLM-R
	try:
	import torch
	import torch.nn as nn
	from transformers import AutoTokenizer, AutoModel
	TORCH_AVAILABLE = True
	except ImportError:
	TORCH_AVAILABLE = False


	# ── HuggingFace Hub — fallback si fichiers absents en local ─────────────────
	HF_REPO_ID = 'BHM03/amane-models'

	def _resolve(local_path: str, hf_filename: str) -> str:
	"""Retourne le chemin local si le fichier existe, sinon télécharge depuis HF Hub."""
	if os.path.exists(local_path):
	return local_path
	try:
	from huggingface_hub import hf_hub_download
	return hf_hub_download(
	repo_id=HF_REPO_ID,
	filename=hf_filename,
	token=os.environ.get('HF_TOKEN'),
	)
	except Exception as e:
	raise FileNotFoundError(
	f'Fichier introuvable localement ({local_path}) '
	f'et téléchargement HF Hub échoué : {e}'
	)


	# ── Chemins modèles ──────────────────────────────────────────────────────────
	BENCHMARK_MODEL_PATH = 'models/ovr_balanced.joblib'
	BENCHMARK_TFIDF_PATH = 'models/tfidf_vectorizer.joblib'
	TEACHER_CKPT_PATH = 'models/teacher/teacher_best.pt'
	TEACHER_TOKENIZER_PATH = 'models/teacher/tokenizer'
	TEACHER_THRESH_PATH = 'models/teacher/thresholds.json'
	TEACHER_STAR_CKPT_PATH = 'models/teacher_star/teacher_star.pt'
	TEACHER_STAR_THRESH_PATH= 'models/teacher_star/threshold_darija.json'
	STUDENT_CKPT_PATH = 'models/student/student_best.pt'
	STUDENT_THRESH_PATH = 'models/student/threshold_student.json'

	# ── Hyperparamètres d'inférence Teacher (cohérents avec l'entraînement) ─────
	TEACHER_MAX_LEN = 256
	TEACHER_STAR_MAX_LEN = 128
	STUDENT_MAX_LEN = 128
	TEACHER_MODEL_HF = 'xlm-roberta-base'
	STUDENT_MODEL_HF = 'distilbert-base-multilingual-cased'


	# ════════════════════════════════════════════════════════════════════════════
	# BENCHMARK — OneVsRest + LogisticRegression
	# ════════════════════════════════════════════════════════════════════════════
	class BenchmarkPredictor:
	def __init__(self):
	self.model = None
	self.tfidf = None
	self.loaded = False
	self.error = None

	def load(self):
	try:
	self.model = joblib.load(_resolve(BENCHMARK_MODEL_PATH, 'ovr_balanced.joblib'))
	self.tfidf = joblib.load(_resolve(BENCHMARK_TFIDF_PATH, 'tfidf_vectorizer.joblib'))
	self.loaded = True
	except Exception as e:
	self.error = str(e)
	self.loaded = False

	def predict(self, text: str) -> dict:
	if not self.loaded:
	raise RuntimeError(f'Benchmark non chargé : {self.error}')
	cleaned = clean_text(text).lower()
	vec = self.tfidf.transform([cleaned])
	probs = self.model.predict_proba(vec)[0]
	return {col: float(round(p, 4)) for col, p in zip(LABEL_COLS, probs)}

	def predict_batch(self, texts: list) -> list[dict]:
	if not self.loaded:
	raise RuntimeError(f'Benchmark non chargé : {self.error}')
	cleaned = [clean_text(t).lower() for t in texts]
	vecs = self.tfidf.transform(cleaned)
	probs = self.model.predict_proba(vecs)
	return [
	{col: float(round(p, 4)) for col, p in zip(LABEL_COLS, row)}
	for row in probs
	]


	# ════════════════════════════════════════════════════════════════════════════
	# XLM-R TEACHER
	# ════════════════════════════════════════════════════════════════════════════
	class TeacherPredictor:
	def __init__(self):
	self.model = None
	self.tokenizer = None
	self.thresholds = None
	self.loaded = False
	self.error = None
	self.device = None

	def load(self):
	if not TORCH_AVAILABLE:
	self.error = (
	'PyTorch / Transformers non installés. '
	'Lance : pip install torch transformers sentencepiece'
	)
	self.loaded = False
	return

	try:
	self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

	# 2. Charger le tokenizer — on utilise directement 'xlm-roberta-base'
	# car le tokenizer ne change pas lors du fine-tuning, et les fichiers
	# locaux sont incomplets (sentencepiece.bpe.model manquant).
	self.tokenizer = AutoTokenizer.from_pretrained(TEACHER_MODEL_HF)

	# 3. Reconstruire l'architecture (identique au notebook d'entraînement)
	class TeacherModel(nn.Module):
	def __init__(self, model_name: str, num_labels: int = 6, dropout: float = 0.1):
	super().__init__()
	self.encoder = AutoModel.from_pretrained(model_name)
	hidden_size = self.encoder.config.hidden_size
	self.dropout = nn.Dropout(dropout)
	self.classifier = nn.Linear(hidden_size, num_labels)

	def mean_pool(self, token_embeds, attention_mask):
	mask = attention_mask.unsqueeze(-1).expand(token_embeds.size()).float()
	summed = (token_embeds * mask).sum(dim=1)
	counts = mask.sum(dim=1).clamp(min=1e-9)
	return summed / counts

	def forward(self, input_ids, attention_mask):
	outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
	pooled = self.mean_pool(outputs.last_hidden_state, attention_mask)
	pooled = self.dropout(pooled)
	return self.classifier(pooled)

	# 4. Charger le checkpoint
	ckpt = torch.load(_resolve(TEACHER_CKPT_PATH, 'teacher_best.pt'), map_location=self.device, weights_only=False)
	model = TeacherModel(TEACHER_MODEL_HF, num_labels=len(LABEL_COLS))
	model.load_state_dict(ckpt['model_state'])
	model.to(self.device)
	model.eval()
	self.model = model

	# 5. Charger les seuils tunés
	if os.path.exists(TEACHER_THRESH_PATH):
	with open(TEACHER_THRESH_PATH) as f:
	self.thresholds = json.load(f)
	else:
	# Fallback : seuil 0.5 par défaut
	self.thresholds = {col: 0.5 for col in LABEL_COLS}

	self.loaded = True

	except Exception as e:
	self.error = f'{type(e).__name__}: {e}'
	self.loaded = False

	# ── Inférence : probabilités ───────────────────────────────────────────
	def predict(self, text: str) -> dict:
	"""Retourne les probabilités sigmoïdes par label."""
	if not self.loaded:
	raise RuntimeError(f'Teacher non chargé : {self.error}')

	cleaned = clean_text(text)
	encoding = self.tokenizer(
	cleaned,
	max_length=TEACHER_MAX_LEN,
	padding='max_length',
	truncation=True,
	return_tensors='pt',
	)
	input_ids = encoding['input_ids'].to(self.device)
	attn_mask = encoding['attention_mask'].to(self.device)

	with torch.no_grad():
	logits = self.model(input_ids, attn_mask)
	probs = torch.sigmoid(logits).cpu().numpy()[0]

	return {col: float(round(p, 4)) for col, p in zip(LABEL_COLS, probs)}

	# ── Inférence : labels binaires (avec seuils tunés) ────────────────────
	def predict_with_labels(self, text: str) -> dict:
	"""
	Retourne probabilités + labels binaires selon les seuils tunés.
	Format : {'probs': {...}, 'labels': {...}, 'thresholds': {...}}
	"""
	probs = self.predict(text)
	labels = {col: int(probs[col] >= self.thresholds[col]) for col in LABEL_COLS}
	return {
	'probs' : probs,
	'labels' : labels,
	'thresholds' : self.thresholds,
	}

	# ── Inférence : batch ──────────────────────────────────────────────────
	def predict_batch(self, texts: list) -> list[dict]:
	"""Inférence batch — plus efficace que predict() en boucle."""
	if not self.loaded:
	raise RuntimeError(f'Teacher non chargé : {self.error}')

	cleaned = [clean_text(t) for t in texts]
	encoding = self.tokenizer(
	cleaned,
	max_length=TEACHER_MAX_LEN,
	padding='max_length',
	truncation=True,
	return_tensors='pt',
	)
	input_ids = encoding['input_ids'].to(self.device)
	attn_mask = encoding['attention_mask'].to(self.device)

	with torch.no_grad():
	logits = self.model(input_ids, attn_mask)
	probs = torch.sigmoid(logits).cpu().numpy()

	return [
	{col: float(round(p, 4)) for col, p in zip(LABEL_COLS, row)}
	for row in probs
	]


	# ════════════════════════════════════════════════════════════════════════════
	# XLM-R TEACHER* — adapté Darija (fine-tuné sur DarLoad, loss masquée)
	# ════════════════════════════════════════════════════════════════════════════
	class TeacherStarPredictor:
	def __init__(self):
	self.model = None
	self.tokenizer = None
	self.threshold = 0.45 # seuil tunée sur DarLoad
	self.loaded = False
	self.error = None
	self.device = None

	def load(self):
	if not TORCH_AVAILABLE:
	self.error = 'PyTorch / Transformers non installés.'
	self.loaded = False
	return

	try:
	self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
	self.tokenizer = AutoTokenizer.from_pretrained(TEACHER_MODEL_HF)

	class TeacherModel(nn.Module):
	def __init__(self, model_name, num_labels=6, dropout=0.1):
	super().__init__()
	self.encoder = AutoModel.from_pretrained(model_name)
	self.dropout = nn.Dropout(dropout)
	self.classifier = nn.Linear(self.encoder.config.hidden_size, num_labels)

	def mean_pool(self, token_embeds, attention_mask):
	mask = attention_mask.unsqueeze(-1).expand(token_embeds.size()).float()
	summed = (token_embeds * mask).sum(dim=1)
	return summed / mask.sum(dim=1).clamp(min=1e-9)

	def forward(self, input_ids, attention_mask):
	out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
	pooled = self.mean_pool(out.last_hidden_state, attention_mask)
	return self.classifier(self.dropout(pooled))

	ckpt = torch.load(_resolve(TEACHER_STAR_CKPT_PATH, 'teacher_star.pt'), map_location=self.device, weights_only=False)
	model = TeacherModel(TEACHER_MODEL_HF, num_labels=len(LABEL_COLS))
	model.load_state_dict(ckpt['model_state'])
	model.to(self.device)
	model.eval()
	self.model = model

	if os.path.exists(TEACHER_STAR_THRESH_PATH):
	with open(TEACHER_STAR_THRESH_PATH) as f:
	self.threshold = json.load(f).get('toxic_darija', 0.45)

	self.loaded = True

	except Exception as e:
	self.error = f'{type(e).__name__}: {e}'
	self.loaded = False

	def predict(self, text: str) -> dict:
	if not self.loaded:
	raise RuntimeError(f'Teacher* non chargé : {self.error}')
	cleaned = clean_text(text)
	encoding = self.tokenizer(
	cleaned, max_length=TEACHER_STAR_MAX_LEN,
	padding='max_length', truncation=True, return_tensors='pt',
	)
	with torch.no_grad():
	logits = self.model(
	encoding['input_ids'].to(self.device),
	encoding['attention_mask'].to(self.device),
	)
	probs = torch.sigmoid(logits).cpu().numpy()[0]
	return {col: float(round(p, 4)) for col, p in zip(LABEL_COLS, probs)}

	def predict_batch(self, texts: list) -> list[dict]:
	if not self.loaded:
	raise RuntimeError(f'Teacher* non chargé : {self.error}')
	cleaned = [clean_text(t) for t in texts]
	encoding = self.tokenizer(
	cleaned, max_length=TEACHER_STAR_MAX_LEN,
	padding='max_length', truncation=True, return_tensors='pt',
	)
	with torch.no_grad():
	logits = self.model(
	encoding['input_ids'].to(self.device),
	encoding['attention_mask'].to(self.device),
	)
	probs = torch.sigmoid(logits).cpu().numpy()
	return [{col: float(round(p, 4)) for col, p in zip(LABEL_COLS, row)} for row in probs]


	# ════════════════════════════════════════════════════════════════════════════
	# FACTORY — chargement unique via st.cache_resource
	# ════════════════════════════════════════════════════════════════════════════
	def load_benchmark() -> BenchmarkPredictor:
	p = BenchmarkPredictor()
	p.load()
	return p

	def load_teacher() -> TeacherPredictor:
	p = TeacherPredictor()
	p.load()
	return p

	def load_teacher_star() -> TeacherStarPredictor:
	p = TeacherStarPredictor()
	p.load()
	return p


	# ════════════════════════════════════════════════════════════════════════════
	# STUDENT — DistilBERT multilingue spécialisé Darija
	# ════════════════════════════════════════════════════════════════════════════
	class StudentPredictor:
	def __init__(self):
	self.model = None
	self.tokenizer = None
	self.threshold = 0.6
	self.loaded = False
	self.error = None
	self.device = None

	def load(self):
	if not TORCH_AVAILABLE:
	self.error = 'PyTorch / Transformers non installés.'
	self.loaded = False
	return

	try:
	self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
	self.tokenizer = AutoTokenizer.from_pretrained(STUDENT_MODEL_HF)

	class StudentModel(nn.Module):
	def __init__(self, model_name, num_labels=6, dropout=0.1):
	super().__init__()
	self.encoder = AutoModel.from_pretrained(model_name)
	self.dropout = nn.Dropout(dropout)
	self.classifier = nn.Linear(self.encoder.config.hidden_size, num_labels)

	def forward(self, input_ids, attention_mask):
	out = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
	pooled = out.last_hidden_state[:, 0, :] # token [CLS]
	return self.classifier(self.dropout(pooled))

	ckpt = torch.load(_resolve(STUDENT_CKPT_PATH, 'student_best.pt'), map_location=self.device, weights_only=False)
	model = StudentModel(STUDENT_MODEL_HF, num_labels=len(LABEL_COLS))
	model.load_state_dict(ckpt['model_state'])
	model.to(self.device)
	model.eval()
	self.model = model

	if os.path.exists(STUDENT_THRESH_PATH):
	with open(STUDENT_THRESH_PATH) as f:
	self.threshold = json.load(f).get('toxic_darija', 0.35)

	self.loaded = True

	except Exception as e:
	self.error = f'{type(e).__name__}: {e}'
	self.loaded = False

	def predict(self, text: str) -> dict:
	if not self.loaded:
	raise RuntimeError(f'Student non chargé : {self.error}')
	cleaned = clean_text(text)
	encoding = self.tokenizer(
	cleaned, max_length=STUDENT_MAX_LEN,
	padding='max_length', truncation=True, return_tensors='pt',
	)
	with torch.no_grad():
	logits = self.model(
	encoding['input_ids'].to(self.device),
	encoding['attention_mask'].to(self.device),
	)
	probs = torch.sigmoid(logits).cpu().numpy()[0]
	return {col: float(round(p, 4)) for col, p in zip(LABEL_COLS, probs)}

	def predict_batch(self, texts: list) -> list[dict]:
	if not self.loaded:
	raise RuntimeError(f'Student non chargé : {self.error}')
	cleaned = [clean_text(t) for t in texts]
	encoding = self.tokenizer(
	cleaned, max_length=STUDENT_MAX_LEN,
	padding='max_length', truncation=True, return_tensors='pt',
	)
	with torch.no_grad():
	logits = self.model(
	encoding['input_ids'].to(self.device),
	encoding['attention_mask'].to(self.device),
	)
	probs = torch.sigmoid(logits).cpu().numpy()
	return [{col: float(round(p, 4)) for col, p in zip(LABEL_COLS, row)} for row in probs]


	def load_student() -> StudentPredictor:
	p = StudentPredictor()
	p.load()
	return p