utils/utils_function.py

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

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"

        if not os.path.exists('utils/data_prepared/'):
            print("Not existed")

            with zipfile.ZipFile('utils/data_prepared.zip', 'r') as zip_ref:
                zip_ref.extractall('utils/')
                
        else:
            print("Existed")

        self.model = TransformerModel.from_pretrained(
            './utils/checkpoints/fon_fr',
            checkpoint_file='checkpoint_best.pt',
            data_name_or_path='utils/data_prepared/',
            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