Trixy7/trixy_core/trainer/core/intent/trainer.py

# -*- coding: utf-8 -*-
"""
Intent-Classifier Trainer.

Trainiert ein Modell fuer Intent-Erkennung.
Exportiert als ONNX fuer schnelle Inferenz auf Raspberry Pi (<100ms).

Architektur:
- Encoder: Sentence-Transformer (multilingual, ~130MB)
- Classifier: 2-Layer Linear (256 hidden, dropout 0.3)
- Export: encoder.onnx + classifier.onnx + tokenizer/

Training auf Desktop (GPU empfohlen), Deployment auf Pi (CPU, ONNX).
"""

from __future__ import annotations

import json
import os
import time
from dataclasses import dataclass, field
from multiprocessing import Event as MPEvent, Queue
from pathlib import Path
from typing import Any

from trixy_core.trainer.itrainer import ITrainer
from trixy_core.trainer.dataclasses import (
    TrainerInfo,
    TrainingProgress,
    ValidationResult,
    FormSchema,
)
from trixy_core.trainer.enums import TrainerType
from trixy_core.trainer.core.intent.slot_lists import SlotRegistry
from trixy_core.trainer.core.intent.data_generator import (
    DataGenerator,
    IntentDefinition,
    TrainingSample,
)
from trixy_core.utils.debug import pinfo, pdebug, perror


class IntentTrainer(ITrainer):
    """
    Trainer fuer Intent-Klassifikation.

    Sammelt Intents aus Plugin-YAML-Dateien, generiert Trainings-Daten
    und trainiert ein ONNX-exportierbares Modell.
    """

    def __init__(self) -> None:
        self._settings: dict[str, Any] = {
            "base_model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
            "samples_per_intent": 200,
            "epochs": 20,
            "batch_size": 32,
            "learning_rate": 2e-5,
            "validation_split": 0.15,
            "slot_lists_file": "config/slot_lists.json",
            "plugins_dir": "plugins",
            "core_nlp_dir": "trixy_core/nlp",
            "output_dir": "models/intent",
            "export_onnx": True,
            "filler_variants": 3,
        }
        self._dataset_config: dict[str, Any] = {
            "additional_examples_dir": "trainer/data/intent/examples",
            "negatives_file": "trainer/data/intent/negatives.txt",
            "processed_dir": "trainer/data/intent/processed",
            "fillers_file": "trainer/assets/intent/fillers_de.json",
            "slot_lists_file": "trainer/assets/intent/slot_lists.json",
        }

    # --- ITrainer Metadaten ---

    @classmethod
    def get_info(cls) -> TrainerInfo:
        return TrainerInfo(
            trainer_id="intent_classifier",
            name="Intent-Classifier",
            description="Trainiert ein Modell fuer Intent-Erkennung und Slot-Extraktion",
            trainer_type=TrainerType.CORE,
            version="1.0.0",
            model_type="intent",
        )

    def get_settings_schema(self) -> FormSchema:
        return FormSchema(fields=[
            {"name": "base_model", "label": "Basis-Modell", "field_type": "select",
             "default": self._settings["base_model"],
             "options": [
                 "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
                 "sentence-transformers/distiluse-base-multilingual-cased-v2",
                 "sentence-transformers/all-MiniLM-L6-v2",
             ],
             "help_text": "Vortrainiertes Encoder-Modell (multilingual = Deutsch)"},
            {"name": "samples_per_intent", "label": "Samples pro Intent", "field_type": "number",
             "default": 200, "min": 50, "max": 2000,
             "help_text": "Anzahl generierter Trainingsbeispiele pro Intent"},
            {"name": "filler_variants", "label": "Fuellwort-Varianten", "field_type": "number",
             "default": 3, "min": 0, "max": 10,
             "help_text": "Varianten pro Beispiel mit Fuellwoertern (0 = deaktiviert)"},
            {"name": "epochs", "label": "Epochen", "field_type": "number",
             "default": 20, "min": 5, "max": 100},
            {"name": "batch_size", "label": "Batch-Groesse", "field_type": "number",
             "default": 32, "min": 8, "max": 128},
            {"name": "learning_rate", "label": "Lernrate", "field_type": "number", "is_float": True,
             "default": 2e-5},
            {"name": "export_onnx", "label": "ONNX Export", "field_type": "checkbox",
             "default": True, "help_text": "Modell als ONNX exportieren (fuer Pi)"},
        ])

    def get_dataset_schema(self) -> FormSchema:
        return FormSchema(fields=[
            {"name": "plugins_dir", "label": "Plugins-Verzeichnis", "field_type": "text",
             "default": "plugins", "help_text": "Verzeichnis mit Plugin-Ordnern"},
            {"name": "additional_examples_dir", "label": "Zusaetzliche Beispiele", "field_type": "text",
             "default": "trainer/data/intent/examples",
             "help_text": "Verzeichnis mit zusaetzlichen Beispiel-Dateien (.txt)"},
            {"name": "negatives_file", "label": "Negativ-Beispiele", "field_type": "text",
             "default": "trainer/data/intent/negatives.txt"},
            {"name": "slot_lists_file", "label": "Slot-Listen", "field_type": "text",
             "default": "config/slot_lists.json"},
        ])

    def get_optional_schema(self) -> FormSchema | None:
        return None

    def get_settings(self) -> dict:
        return dict(self._settings)

    def set_settings(self, settings: dict) -> None:
        self._settings.update(settings)

    def get_dataset_info(self) -> dict:
        """Analysiert das vorhandene Dataset."""
        gen = self._create_generator()
        gen.collect_from_yaml_files(
            self._settings.get("plugins_dir", "plugins"),
            self._settings.get("core_nlp_dir", "trixy_core/nlp"),
        )
        stats = gen.get_stats()

        # Bestehendes Modell pruefen
        model_dir = Path(self._settings.get("output_dir", "models/intent"))
        metafile = model_dir / "metafile.json"
        existing_model = None
        if metafile.is_file():
            try:
                with open(metafile) as f:
                    existing_model = json.load(f)
            except Exception:
                pass

        return {
            **stats,
            "existing_model": existing_model,
        }

    def set_dataset_config(self, config: dict) -> None:
        self._dataset_config.update(config)
        if "plugins_dir" in config:
            self._settings["plugins_dir"] = config["plugins_dir"]

    def validate(self) -> ValidationResult:
        """Validiert die Trainings-Konfiguration."""
        errors = []

        # Abhaengigkeiten pruefen
        try:
            import torch
        except ImportError:
            errors.append("PyTorch nicht installiert (pip install torch)")

        try:
            from sentence_transformers import SentenceTransformer
        except ImportError:
            errors.append("sentence-transformers nicht installiert")

        try:
            import sklearn
        except ImportError:
            errors.append("scikit-learn nicht installiert")

        if errors:
            return ValidationResult(valid=False, errors=errors)

        # Intents pruefen
        gen = self._create_generator()
        count = gen.collect_from_yaml_files(
            self._settings.get("plugins_dir", "plugins"),
            self._settings.get("core_nlp_dir", "trixy_core/nlp"),
        )

        if count == 0:
            return ValidationResult(
                valid=False,
                errors=["Keine Intents gefunden — YAML-Dateien pruefen"],
            )

        return ValidationResult(valid=True)

    # --- Training ---

    def train(
        self,
        progress_queue: Queue,
        pause_flag: MPEvent,
        stop_flag: MPEvent,
    ) -> None:
        """
        Fuehrt das Training durch (laeuft im Subprocess).

        Sendet TrainingProgress ueber die Queue.
        Prueft stop_flag und pause_flag in Schleifen.
        """
        import torch
        import numpy as np
        from sklearn.model_selection import StratifiedKFold
        from sklearn.metrics import classification_report
        from sklearn.preprocessing import LabelEncoder

        total_epochs = self._settings.get("epochs", 20)

        def report(
            progress: float, message: str,
            epoch: int = 0, loss: float = 0.0, val_loss: float = 0.0, accuracy: float = 0.0,
        ) -> None:
            from trixy_core.trainer.enums import TrainingState
            progress_queue.put(TrainingProgress(
                state=TrainingState.TRAINING,
                epoch=epoch,
                total_epochs=total_epochs,
                loss=loss,
                val_loss=val_loss,
                accuracy=accuracy,
                message=message,
            ))

        def check_stop() -> bool:
            return stop_flag.is_set()

        def wait_if_paused() -> None:
            while pause_flag.is_set() and not stop_flag.is_set():
                time.sleep(0.5)

        report(0.0, "Initialisiere...")

        # Settings-Typen sicherstellen (TUI sendet alles als String)
        for int_key in ("samples_per_intent", "epochs", "batch_size"):
            if int_key in self._settings:
                try:
                    self._settings[int_key] = int(self._settings[int_key])
                except (ValueError, TypeError):
                    pass
        for float_key in ("learning_rate", "validation_split"):
            if float_key in self._settings:
                try:
                    self._settings[float_key] = float(self._settings[float_key])
                except (ValueError, TypeError):
                    pass
        if "export_onnx" in self._settings:
            val = self._settings["export_onnx"]
            if isinstance(val, str):
                self._settings["export_onnx"] = val.lower() in ("true", "1", "yes")

        # === 1. Slot-Listen laden ===
        report(0.05, "Lade Slot-Listen...")
        slot_registry = SlotRegistry()

        # Slot-Listen: trainer/assets/intent/ bevorzugt, Fallback config/
        slot_file = self._dataset_config.get(
            "slot_lists_file", "trainer/assets/intent/slot_lists.json",
        )
        if not os.path.isfile(slot_file):
            slot_file = "config/slot_lists.json"  # Fallback
        if os.path.isfile(slot_file):
            slot_registry.load_from_file(slot_file)

        # Plugin-eigene Slot-Listen laden (plugins/*/slot_lists.json)
        plugins_dir = self._settings.get("plugins_dir", "plugins")
        if os.path.isdir(plugins_dir):
            from pathlib import Path
            for plugin_slots in sorted(Path(plugins_dir).glob("*/slot_lists.json")):
                loaded = slot_registry.load_from_file(plugin_slots)
                if loaded:
                    report(0.05, f"Plugin-Slots: {plugin_slots.parent.name} ({loaded})")

        # === 2. Intents sammeln ===
        report(0.1, "Sammle Intent-Definitionen...")
        generator = self._create_generator(slot_registry)
        generator.collect_from_yaml_files(
            self._settings.get("plugins_dir", "plugins"),
            self._settings.get("core_nlp_dir", "trixy_core/nlp"),
        )

        # Fuellwoerter laden: trainer/assets/intent/ bevorzugt, Fallback config/
        fillers_file = self._dataset_config.get(
            "fillers_file", "trainer/assets/intent/fillers_de.json",
        )
        if not os.path.isfile(fillers_file):
            fillers_file = "config/fillers_de.json"
        generator.load_fillers(fillers_file)

        # Zusaetzliche Beispiele laden
        examples_dir = self._dataset_config.get("additional_examples_dir", "")
        if examples_dir and os.path.isdir(examples_dir):
            self._load_additional_examples(generator, examples_dir)

        if generator.intent_count == 0:
            report(1.0, "FEHLER: Keine Intents gefunden")
            return

        if check_stop():
            return

        # === 3. Trainings-Daten generieren ===
        report(0.15, f"Generiere Trainings-Daten ({generator.intent_count} Intents)...")
        samples_per_intent = self._settings.get("samples_per_intent", 200)
        filler_variants = self._settings.get("filler_variants", 3)
        samples = generator.generate(
            samples_per_intent=samples_per_intent,
            filler_variants=filler_variants,
        )

        # Trainings-Daten als lesbare Dateien exportieren
        processed_dir = self._dataset_config.get(
            "processed_dir", "trainer/data/intent/processed",
        )
        report(0.17, "Exportiere Trainings-Daten...")
        generator.export_to_files(samples, processed_dir)

        # === 4. Label-Encoding ===
        report(0.2, "Encode Labels...")
        intent_encoder = LabelEncoder()
        all_intents = [s.intent for s in samples]
        intent_labels = intent_encoder.fit_transform(all_intents)
        intent_names = list(intent_encoder.classes_)
        num_intents = len(intent_names)

        pinfo(f"Intent-Klassen: {num_intents}, Samples: {len(samples)}")

        # === 5. Sentence-Encoder laden ===
        report(0.25, "Lade Sentence-Encoder...")
        base_model = self._settings.get(
            "base_model",
            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
        )

        try:
            from sentence_transformers import SentenceTransformer
            encoder = SentenceTransformer(base_model)
            embedding_dim = encoder.get_sentence_embedding_dimension()
            pinfo(f"Encoder: {base_model} (dim={embedding_dim})")
        except Exception as e:
            report(1.0, f"FEHLER: Encoder laden fehlgeschlagen: {e}")
            return

        if check_stop():
            return

        # === 6. Embeddings berechnen (mit Cache) ===
        report(0.3, f"Berechne Embeddings ({len(samples)} Samples)...")
        texts = [s.text for s in samples]

        from trixy_core.trainer.core.intent.embedding_cache import EmbeddingCache
        model_dir = Path(self._settings.get("output_dir", "models/intent"))
        cache = EmbeddingCache(model_dir)
        cache.load(base_model)

        def cache_progress(computed: int, total: int, cached: int) -> None:
            if total > 0:
                pct = computed / total
                report(0.3 + pct * 0.1, f"Embeddings: {cached} Cache, {computed}/{total} neu")

        embeddings = cache.get_or_compute(
            texts, encoder.encode, batch_size=64,
            progress_fn=cache_progress,
        )
        cache.save()

        if check_stop():
            return

        # === 7. Cross-Validation Training ===
        report(0.35, "Starte Cross-Validation Training...")
        epochs = self._settings.get("epochs", 20)
        batch_size = self._settings.get("batch_size", 32)
        lr = self._settings.get("learning_rate", 2e-5)

        n_splits = min(5, min(
            np.bincount(intent_labels).min(),
            len(intent_labels) // 5,
        ))
        n_splits = max(2, n_splits)

        skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)
        fold_results: list[dict] = []
        best_model_state = None
        best_val_acc = 0.0
        all_train_losses: list[float] = []
        all_val_losses: list[float] = []

        for fold, (train_idx, val_idx) in enumerate(skf.split(embeddings, intent_labels)):
            if check_stop():
                return
            wait_if_paused()

            report(
                0.35 + (fold / n_splits) * 0.45,
                f"Fold {fold + 1}/{n_splits}...",
            )

            X_train = torch.FloatTensor(embeddings[train_idx])
            y_train = torch.LongTensor(intent_labels[train_idx])
            X_val = torch.FloatTensor(embeddings[val_idx])
            y_val = torch.LongTensor(intent_labels[val_idx])

            # Classifier erstellen
            classifier = torch.nn.Sequential(
                torch.nn.Linear(embedding_dim, 256),
                torch.nn.ReLU(),
                torch.nn.Dropout(0.3),
                torch.nn.Linear(256, num_intents),
            )
            optimizer = torch.optim.Adam(classifier.parameters(), lr=lr)
            loss_fn = torch.nn.CrossEntropyLoss()

            for epoch in range(epochs):
                if check_stop():
                    return
                wait_if_paused()

                classifier.train()
                epoch_loss = 0.0

                indices = torch.randperm(len(X_train))
                for i in range(0, len(indices), batch_size):
                    batch_idx = indices[i:i + batch_size]
                    output = classifier(X_train[batch_idx])
                    loss = loss_fn(output, y_train[batch_idx])
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()
                    epoch_loss += loss.item()

                avg_loss = epoch_loss / max(1, len(X_train) // batch_size)

                # Validierung
                classifier.eval()
                with torch.no_grad():
                    val_output = classifier(X_val)
                    val_loss = loss_fn(val_output, y_val).item()
                    val_pred = val_output.argmax(dim=1)
                    val_acc = (val_pred == y_val).float().mean().item()

                if fold == 0:
                    all_train_losses.append(avg_loss)
                    all_val_losses.append(val_loss)

                progress = 0.35 + ((fold + epoch / epochs) / n_splits) * 0.45
                report(progress, f"Fold {fold+1} Ep {epoch+1}/{epochs} — Acc: {val_acc:.3f}", avg_loss)

            # Fold-Ergebnis
            classifier.eval()
            with torch.no_grad():
                val_output = classifier(X_val)
                val_pred = val_output.argmax(dim=1).numpy()
                fold_acc = (val_pred == y_val.numpy()).mean()

            fold_results.append({"fold": fold + 1, "accuracy": float(fold_acc)})

            if fold_acc > best_val_acc:
                best_val_acc = fold_acc
                best_model_state = classifier.state_dict()

        # === 8. Per-Intent Metriken ===
        report(0.82, "Berechne Metriken...")

        # Bestes Modell auf allen Daten evaluieren
        classifier = torch.nn.Sequential(
            torch.nn.Linear(embedding_dim, 256),
            torch.nn.ReLU(),
            torch.nn.Dropout(0.3),
            torch.nn.Linear(256, num_intents),
        )
        classifier.load_state_dict(best_model_state)
        classifier.eval()

        with torch.no_grad():
            all_output = classifier(torch.FloatTensor(embeddings))
            all_pred = all_output.argmax(dim=1).numpy()

        report_str = classification_report(
            intent_labels, all_pred,
            target_names=intent_names,
            output_dict=False,
        )
        report_dict = classification_report(
            intent_labels, all_pred,
            target_names=intent_names,
            output_dict=True,
        )
        pinfo(f"Classification Report:\n{report_str}")

        # === 9. Modell speichern ===
        report(0.88, "Speichere Modell...")
        model_dir = Path(self._settings.get("output_dir", "models/intent"))
        model_dir.mkdir(parents=True, exist_ok=True)

        torch.save(best_model_state, model_dir / "classifier.pth")

        metafile = {
            "model_name": "intent_classifier",
            "base_model": base_model,
            "embedding_dim": embedding_dim,
            "intent_names": intent_names,
            "slot_names": list(generator._slots.slot_names),
            "num_intents": num_intents,
            "train_samples": len(samples),
            "best_val_accuracy": float(best_val_acc),
            "cross_validation_folds": fold_results,
            "mean_cv_accuracy": float(np.mean([f["accuracy"] for f in fold_results])),
            "epochs": epochs,
            "train_losses": all_train_losses,
            "val_losses": all_val_losses,
            "per_intent_metrics": {
                name: report_dict.get(name, {})
                for name in intent_names
                if name != "__negative__"
            },
            "created": time.strftime("%Y-%m-%d %H:%M:%S"),
        }
        with open(model_dir / "metafile.json", "w") as f:
            json.dump(metafile, f, indent=2, ensure_ascii=False)

        with open(model_dir / "intent_labels.json", "w") as f:
            json.dump(intent_names, f, ensure_ascii=False)

        # === 10. BIO Slot-Tagger trainieren ===
        if check_stop():
            return

        report(0.88, "Trainiere BIO Slot-Tagger...")
        try:
            self._train_bio_tagger(
                samples, encoder, model_dir, epochs,
                batch_size, report, check_stop, wait_if_paused,
            )
        except Exception as e:
            perror(f"BIO-Tagger Training fehlgeschlagen: {e}")
            import traceback
            pdebug(traceback.format_exc())

        # === 11. ONNX Export ===
        if self._settings.get("export_onnx", True):
            report(0.92, "Exportiere Classifier nach ONNX...")
            try:
                dummy_input = torch.randn(1, embedding_dim)
                torch.onnx.export(
                    classifier,
                    dummy_input,
                    str(model_dir / "intent_classifier.onnx"),
                    input_names=["embedding"],
                    output_names=["intent_logits"],
                    dynamic_axes={
                        "embedding": {0: "batch"},
                        "intent_logits": {0: "batch"},
                    },
                    opset_version=18,
                )
                pinfo(f"ONNX Classifier: {model_dir / 'intent_classifier.onnx'}")
            except Exception as e:
                perror(f"ONNX Classifier Export fehlgeschlagen: {e}")

            # Encoder ONNX Export
            report(0.95, "Exportiere Encoder nach ONNX...")
            try:
                self._export_encoder_onnx(encoder, model_dir)
            except Exception as e:
                perror(f"ONNX Encoder Export fehlgeschlagen: {e}")
                pinfo("Encoder muss zur Laufzeit via SentenceTransformer geladen werden")

            # Tokenizer speichern
            try:
                tokenizer = encoder.tokenizer
                tokenizer_dir = model_dir / "tokenizer"
                tokenizer.save_pretrained(str(tokenizer_dir))
                pinfo(f"Tokenizer gespeichert: {tokenizer_dir}")
            except Exception as e:
                pdebug(f"Tokenizer speichern fehlgeschlagen: {e}")

        report(1.0, f"Training abgeschlossen! Val-Acc: {best_val_acc:.3f}")
        pinfo(
            f"Intent-Classifier Training abgeschlossen: "
            f"{num_intents} Intents, Val-Acc={best_val_acc:.3f}, "
            f"CV-Mean={np.mean([f['accuracy'] for f in fold_results]):.3f}"
        )

    # === BIO Slot-Tagger ===

    def _train_bio_tagger(
        self,
        samples: list,
        encoder: Any,
        model_dir: Path,
        epochs: int,
        batch_size: int,
        report: Any,
        check_stop: Any,
        wait_if_paused: Any,
    ) -> None:
        """
        Trainiert den BIO-Slot-Tagger auf Token-Level Embeddings.

        Verwendet den gleichen Encoder wie der Intent-Classifier.
        Nur Samples mit Slots werden fuer das BIO-Training genutzt.
        """
        import torch
        import numpy as np
        from trixy_core.trainer.core.intent.bio_tagger import (
            BIOTagSet, generate_word_bio_tags, align_bio_to_tokens,
        )

        MAX_SEQ_LEN = 64

        # Tokenizer holen
        tokenizer = encoder.tokenizer

        # Samples mit Slots filtern
        slot_samples = [s for s in samples if s.slots and s.intent != "__negative__"]
        if not slot_samples:
            pinfo("BIO-Tagger: Keine Samples mit Slots — uebersprungen")
            return

        # Alle Slot-Namen sammeln
        all_slot_names: set[str] = set()
        for s in slot_samples:
            all_slot_names.update(s.slots.keys())

        # Auch Samples ohne Slots als Negative nehmen (lernt "O" Tags)
        no_slot_samples = [s for s in samples if not s.slots and s.intent != "__negative__"]
        import random
        rng = random.Random(42)
        neg_count = min(len(no_slot_samples), len(slot_samples))
        if no_slot_samples:
            neg_bio_samples = rng.sample(no_slot_samples, neg_count)
        else:
            neg_bio_samples = []

        bio_samples = slot_samples + neg_bio_samples
        rng.shuffle(bio_samples)

        pinfo(
            f"BIO-Tagger: {len(slot_samples)} Slot-Samples, "
            f"{len(neg_bio_samples)} Neg-Samples, "
            f"{len(all_slot_names)} Slot-Typen: {sorted(all_slot_names)}"
        )

        # BIO-TagSet erstellen
        tagset = BIOTagSet(slot_names=sorted(all_slot_names))
        pinfo(f"BIO-Tags: {tagset.num_tags} ({tagset.tags})")

        # Token-Embeddings und BIO-Labels berechnen
        report(0.89, f"BIO: Berechne Token-Embeddings ({len(bio_samples)} Samples)...")

        # Transformer-Modell aus SentenceTransformer extrahieren
        transformer = encoder[0]
        auto_model = transformer.auto_model.cpu()  # Auf CPU (vermeidet CUDA-Mismatch)
        auto_model.eval()

        all_token_embs: list[np.ndarray] = []
        all_bio_labels: list[np.ndarray] = []

        for i, sample in enumerate(bio_samples):
            if check_stop():
                return

            # 1. Wort-Level BIO-Tags
            word_tags = generate_word_bio_tags(sample.text, sample.slots)

            # 2. Tokenisieren
            encoding = tokenizer(
                sample.text,
                padding="max_length",
                truncation=True,
                max_length=MAX_SEQ_LEN,
                return_tensors="pt",
                return_offsets_mapping=True,
            )

            # 3. Token-Level BIO-Tags (Subword-aligned)
            token_tags = align_bio_to_tokens(
                word_tags, sample.text, tokenizer, MAX_SEQ_LEN,
            )

            # 4. BIO-Label Indizes
            label_indices = [tagset.tag_to_idx(t) for t in token_tags]
            all_bio_labels.append(np.array(label_indices, dtype=np.int64))

            # 5. Token-Embeddings via Encoder
            with torch.no_grad():
                input_ids = encoding["input_ids"]
                attention_mask = encoding["attention_mask"]
                outputs = auto_model(input_ids=input_ids, attention_mask=attention_mask)
                token_emb = outputs.last_hidden_state[0].cpu().numpy()  # (seq_len, 384)

            all_token_embs.append(token_emb)

            if (i + 1) % 500 == 0:
                report(0.89 + (i / len(bio_samples)) * 0.03,
                       f"BIO: Token-Embeddings {i+1}/{len(bio_samples)}")

        # Zu Tensoren
        X_bio = torch.FloatTensor(np.stack(all_token_embs))   # (N, seq_len, 384)
        y_bio = torch.LongTensor(np.stack(all_bio_labels))    # (N, seq_len)

        embedding_dim = X_bio.shape[2]

        # BIO-Kopf erstellen
        bio_head = torch.nn.Sequential(
            torch.nn.Linear(embedding_dim, 128),
            torch.nn.ReLU(),
            torch.nn.Dropout(0.3),
            torch.nn.Linear(128, tagset.num_tags),
        )

        optimizer = torch.optim.Adam(bio_head.parameters(), lr=1e-3)
        loss_fn = torch.nn.CrossEntropyLoss(ignore_index=tagset.pad_idx)

        # Training
        bio_epochs = min(epochs, 20)
        report(0.92, f"BIO: Trainiere {bio_epochs} Epochen...")

        for epoch in range(bio_epochs):
            if check_stop():
                return
            wait_if_paused()

            bio_head.train()
            epoch_loss = 0.0

            indices = torch.randperm(len(X_bio))
            for i in range(0, len(indices), batch_size):
                batch_idx = indices[i:i + batch_size]
                batch_x = X_bio[batch_idx]    # (B, seq_len, 384)
                batch_y = y_bio[batch_idx]    # (B, seq_len)

                # Forward: (B, seq_len, 384) → (B, seq_len, num_tags)
                logits = bio_head(batch_x)

                # Loss: Reshape fuer CrossEntropy
                # (B*seq_len, num_tags) vs (B*seq_len,)
                loss = loss_fn(
                    logits.view(-1, tagset.num_tags),
                    batch_y.view(-1),
                )

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                epoch_loss += loss.item()

            avg_loss = epoch_loss / max(1, len(X_bio) // batch_size)

            if (epoch + 1) % 5 == 0 or epoch == 0:
                # Accuracy (ohne PAD)
                bio_head.eval()
                with torch.no_grad():
                    all_logits = bio_head(X_bio)
                    preds = all_logits.argmax(dim=2)  # (N, seq_len)

                    mask = y_bio != tagset.pad_idx
                    correct = ((preds == y_bio) & mask).sum().item()
                    total = mask.sum().item()
                    acc = correct / total if total > 0 else 0

                    # Slot-spezifische Accuracy (nicht-O Tags)
                    slot_mask = mask & (y_bio != tagset.o_idx)
                    slot_correct = ((preds == y_bio) & slot_mask).sum().item()
                    slot_total = slot_mask.sum().item()
                    slot_acc = slot_correct / slot_total if slot_total > 0 else 0

                pinfo(
                    f"BIO Epoch {epoch+1}/{bio_epochs}: "
                    f"loss={avg_loss:.4f}, acc={acc:.3f}, slot_acc={slot_acc:.3f}"
                )

        # Speichern
        report(0.95, "BIO: Speichere Modell...")
        torch.save(bio_head.state_dict(), model_dir / "slot_tagger.pth")

        # BIO-Labels speichern
        with open(model_dir / "bio_labels.json", "w") as f:
            json.dump(tagset.to_dict(), f, indent=2, ensure_ascii=False)

        # ONNX Export des BIO-Kopfs
        try:
            bio_head.eval()
            dummy = torch.randn(1, MAX_SEQ_LEN, embedding_dim)
            torch.onnx.export(
                bio_head,
                dummy,
                str(model_dir / "slot_tagger.onnx"),
                input_names=["token_embeddings"],
                output_names=["bio_logits"],
                dynamic_axes={
                    "token_embeddings": {0: "batch", 1: "seq_len"},
                    "bio_logits": {0: "batch", 1: "seq_len"},
                },
                opset_version=18,
            )
            size_kb = (model_dir / "slot_tagger.onnx").stat().st_size / 1024
            pinfo(f"BIO ONNX: slot_tagger.onnx ({size_kb:.1f} KB)")
        except Exception as e:
            perror(f"BIO ONNX Export fehlgeschlagen: {e}")

        pinfo(
            f"BIO-Tagger Training abgeschlossen: "
            f"{tagset.num_tags} Tags, {len(bio_samples)} Samples"
        )

    # === Hilfsmethoden ===

    def _create_generator(self, slot_registry: SlotRegistry | None = None) -> DataGenerator:
        """Erstellt einen DataGenerator mit konfigurierter SlotRegistry."""
        if slot_registry is None:
            slot_registry = SlotRegistry()
            slot_file = self._settings.get("slot_lists_file", "")
            if slot_file and os.path.isfile(slot_file):
                slot_registry.load_from_file(slot_file)

        gen = DataGenerator(slot_registry)
        neg_file = self._dataset_config.get("negatives_file", "trainer/data/intent/negatives.txt")
        if neg_file:
            gen._negatives_file = neg_file
        return gen

    def _load_additional_examples(
        self, generator: DataGenerator, examples_dir: str,
    ) -> None:
        """
        Laedt zusaetzliche Beispiele aus .txt Dateien.

        Dateiformat:
        ```
        # intent: weather_forecast
        wie wird das wetter morgen
        wird es regnen
        ```
        """
        for txt_file in sorted(Path(examples_dir).glob("*.txt")):
            intent_name = ""
            examples = []

            with open(txt_file, encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if line.startswith("# intent:"):
                        intent_name = line.split(":", 1)[1].strip()
                    elif line and not line.startswith("#"):
                        examples.append(line)

            if intent_name and examples:
                generator.add_intent(IntentDefinition(
                    name=intent_name,
                    examples=examples,
                ))
                pdebug(f"Beispiele: {intent_name} ({len(examples)} aus {txt_file.name})")

    def _export_encoder_onnx(
        self, encoder: Any, model_dir: Path,
    ) -> None:
        """
        Exportiert den SentenceTransformer-Encoder nach ONNX.

        Nutzt die optimum-Library fuer zuverlaessigen Export.
        Quantisiert auf Int8 (~112MB statt 448MB).
        Auf dem Pi braucht man dann kein PyTorch/sentence-transformers.
        """
        try:
            from optimum.onnxruntime import ORTModelForFeatureExtraction, ORTQuantizer
            from optimum.onnxruntime.configuration import AutoQuantizationConfig
            from transformers import AutoTokenizer
        except ImportError:
            perror("optimum nicht installiert — pip install optimum[onnxruntime]")
            return

        base_model = self._settings.get(
            "base_model",
            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
        )

        encoder_dir = model_dir / "encoder_onnx"
        encoder_float_dir = model_dir / "encoder_float32"

        # 1. Float32 ONNX exportieren
        pinfo("ONNX Encoder: Exportiere float32...")
        model = ORTModelForFeatureExtraction.from_pretrained(base_model, export=True)
        model.save_pretrained(str(encoder_float_dir))

        # 2. Tokenizer speichern
        tokenizer = AutoTokenizer.from_pretrained(base_model)
        tokenizer.save_pretrained(str(encoder_float_dir))

        # 3. Quantisierung (Int8 — ~4x kleiner)
        pinfo("ONNX Encoder: Quantisiere auf Int8...")
        quantizer = ORTQuantizer.from_pretrained(str(encoder_float_dir))
        qconfig = AutoQuantizationConfig.avx2(is_static=False)
        quantizer.quantize(save_dir=str(encoder_dir), quantization_config=qconfig)

        # Tokenizer ins quantisierte Verzeichnis kopieren
        tokenizer.save_pretrained(str(encoder_dir))

        # Float32-Version loeschen (spart ~340MB)
        import shutil
        shutil.rmtree(str(encoder_float_dir), ignore_errors=True)

        # Groesse anzeigen
        total_size = sum(
            f.stat().st_size for f in encoder_dir.rglob("*") if f.is_file()
        )
        pinfo(f"ONNX Encoder: {encoder_dir} ({total_size / 1024 / 1024:.1f} MB, Int8)")