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

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(
            f'./utils/checkpoints/{inner}',
            checkpoint_file = 'checkpoint_best.pt',
            data_name_or_path = f'utils/datas/data_prepared_{inner}/',
            source_lang='fon' if isFon else 'fr',
            target_lang='fr' if isFon else 'fon'
        )

        # 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)
        
        # Utilisation de la méthode generate avec le paramètre beam
        translations = self.model.generate(tokens, beam=5)
        print(type(translations))
        print(translations[0])
        best_translation_tokens = [translations[i]['tokens'].tolist() for i in range(5)]

        # Décodage des tokens en traduction
        translations = [self.model.decode(best_translation_tokens[i]) for i in range(5)]
        
        return "\n".join(translations)