from fairseq.models.transformer import TransformerModel
import os
import torch

class Translator:
    def __init__(self, isFon:bool, device='cuda' if torch.cuda.is_available() else 'cpu'):

        # Charger le modèle pré-entraîné avec Fairseq
        inner = "fon_fr" if isFon else "fr_fon"

        self.model = TransformerModel.from_pretrained(
            './utils/checkpoints/fon_fr',
            #utils/checkpoints/fon_fr
            checkpoint_file='checkpoint_best.pt',  # Nom du fichier de point de contrôle du modèle
            data_name_or_path='utils/data_prepared/',  # Chemin vers le dossier contenant les données du modèle
            source_lang='fon',
            target_lang='fr'
        )

        print("#########################")
        print(type(self.model))
        print("#########################")
        # Définir le périphérique sur lequel exécuter le modèle (par défaut sur 'cuda' si disponible)
        self.model.to(device)
        
        # Mettre le modèle en mode évaluation (pas de mise à jour des poids)
        self.model.eval()

    def translate(self, text):
        # Encodage du texte en tokens
        tokens = self.model.encode(text)
        
        # Décodage des tokens en traduction
        translation = self.model.decode(tokens)
        
        return translation