Spaces:
Sleeping
Sleeping
| 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) | |