File manager - Edit - /usr/local/lib/python3.9/dist-packages/pythainlp/translate/small100.py
Back
from transformers import M2M100ForConditionalGeneration from .tokenization_small100 import SMALL100Tokenizer class Small100Translator: """ Machine Translation using small100 model - Huggingface https://huggingface.co/alirezamsh/small100 :param bool use_gpu : load model using GPU (Default is False) """ def __init__( self, use_gpu: bool = False, pretrained: str = "alirezamsh/small100", ) -> None: self.pretrained = pretrained self.model = M2M100ForConditionalGeneration.from_pretrained(self.pretrained) self.tgt_lang = None if use_gpu: self.model = self.model.cuda() def translate(self, text: str, tgt_lang: str="en") -> str: """ Translate text from X to X :param str text: input text in source language :param str tgt_lang: target language :return: translated text in target language :rtype: str :Example: :: from pythainlp.translate.small100 import Small100Translator mt = Small100Translator() # Translate text from Thai to English mt.translate("ทดสอบระบบ", tgt_lang="en") # output: 'Testing system' # Translate text from Thai to Chinese mt.translate("ทดสอบระบบ", tgt_lang="zh") # output: '系统测试' # Translate text from Thai to French mt.translate("ทดสอบระบบ", tgt_lang="fr") # output: 'Test du système' """ if tgt_lang!=self.tgt_lang: self.tokenizer = SMALL100Tokenizer.from_pretrained(self.pretrained, tgt_lang=tgt_lang) self.tgt_lang = tgt_lang self.translated = self.model.generate( **self.tokenizer(text, return_tensors="pt") ) return self.tokenizer.batch_decode(self.translated, skip_special_tokens=True)[0]
| ver. 1.4 |
Github
|
.
| PHP 7.4.33 | Generation time: 0.76 |
proxy
|
phpinfo
|
Settings