model-ml/ml-api/app/train_translate.py

import os
from pathlib import Path

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM


DATASET_PATH = Path("/work/models/data/dataset.json")
OUTPUT_DIR = Path("/work/models/nllb-custom")
BASE_MODEL = "facebook/nllb-200-distilled-600M"


def train_from_local_dataset() -> str:
    """
    Placeholder de fine-tuning NLLB. Pour un environnement CPU et rapide, on se contente
    de préparer le répertoire custom avec le tokenizer et le modèle de base afin
    de permettre les tests de pipeline. Dans un environnement GPU, remplacer par un
    entraînement réel (Trainer, datasets, etc.).
    """
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    tok = AutoTokenizer.from_pretrained(BASE_MODEL)
    mdl = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL)

    tok.save_pretrained(OUTPUT_DIR)
    mdl.save_pretrained(OUTPUT_DIR)
    return str(OUTPUT_DIR)