Source code for pororo.tasks.word_translation

"""Word Translation related modeling class"""

from collections import namedtuple
from typing import Optional

from pororo.tasks.utils.base import PororoFactoryBase, PororoSimpleBase


[docs]class PororoWordTranslationFactory(PororoFactoryBase): """ Conduct cross-word translation tasks See also: word2word: A Collection of Bilingual Lexicons for 3,564 Language Pairs (https://arxiv.org/abs/1911.12019) Multi (`word2word`) - dataset: word2word (Choe et al. 2019) - metric: N/A Args: word (str): input word to be translated Returns: List[str]: word translation candidates Examples: >>> wt = Pororo(task="word_translation", lang="en", tgt="fr") >>> wt("apple") ['pomme', 'pommier', 'pommes', 'tombe', 'yeux'] >>> wt = Pororo(task="word_translation", lang="ja", tgt="ko") >>> wt("リンゴ") ['선악과', '사과', '에덴', '링고', '귀하'] >>> wt = Pororo(task="word_translation", lang="ko", tgt="en") >>> wt("사과") ['apologize', 'apology', 'apple', 'apologies', 'apologizing'] """ # TODO: How to follow other module's __init__ concept? def __init__(self, task: str, lang: str, model: Optional[str], tgt: str): if model is None: model = "word2word" assert model in ["word2word"] Config = namedtuple("Config", ["task", "lang", "n_model"]) self.config = Config(task, lang, model) self._tgt = tgt
[docs] def load(self, device: str): """ Load user-selected task-specific model Args: device (str): device information Returns: object: User-selected task-specific model """ from word2word import Word2word model = Word2word(self.config.lang, self._tgt) return PororoWord2Word(model, self.config)
[docs]class PororoWord2Word(PororoSimpleBase): def __init__(self, model, config): super().__init__(config) self._model = model
[docs] def predict(self, word: str, **kwargs): """ Conduct word-to-word translation Args: word (str): input word to be translated Returns: List[str]: word translation candidates """ return self._model(word)