Spaces:
No application file
No application file
File size: 3,568 Bytes
888d016 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
from lingua import LanguageDetectorBuilder, Language
class Translator:
def __init__(self, languages:list=None, model_size:str='418M'):
"""Detects and translates text into a required language, using the
M2M100 model and the Lingua package. If the language is being detected
from a pool of possible languages these can be stated to improve
computational efficiency, otherwise leave blank to translate from any
language.
Args:
languages (list, optional): A list of potential source languages as
ISO-639-1 codes. Leave as None if source language is unknown.
Defaults to None.
model_str (str, optional): The model being used. Can be '418M' or
'1.2B'. Defaults to '418M'.
"""
if languages:
self.languages = [getattr(Language, l.upper()) for l in languages]
else:
self.languages = None
self.detector = self.get_detector()
self.model_str = f'facebook/m2m100_{model_size}'
self.model = M2M100ForConditionalGeneration.from_pretrained(self.model_str)
def get_detector(self)-> LanguageDetectorBuilder:
"""Retrieves the language detection model. If a list of potential
languages has been provided in the class initialisation then the
detector will chose from those classes.
Returns:
LanguageDetectorBuilder: initialised laguage detection model.
"""
if self.languages:
detector = LanguageDetectorBuilder.from_iso_codes_639_1(*self.languages)
else:
detector = LanguageDetectorBuilder.from_all_languages()
return detector.build()
def translate(self, text:str, out_lang:str)->str:
"""translates text to the language defined by out_lang. Source language
is detected automatically.
Args:
text (str): text to be translated
out_lang (str): ISO Code 639-1 of target language (e.g. "en")
Returns:
str: translated text in out_lang
"""
src_lang = self.detect_language(text)
src_tokenizer = self.get_tokenizer(src_lang)
src_tokens = src_tokenizer(text, return_tensors='pt')
out_tokens = self.model.generate(**src_tokens, forced_bos_token_id=src_tokenizer.get_lang_id(out_lang))
out_text = src_tokenizer.batch_decode(out_tokens, skip_special_tokens=True)
return {'lanuage':src_lang, 'translation':out_text}
def get_tokenizer(self, src_lang:str)->M2M100Tokenizer:
"""Retrieves the tokenizer in the required source language. If the
Args:
src_lang (str): ISO0-639-1 country code
Returns:
M2M100Tokenizer: _description_
"""
try:
return M2M100Tokenizer.from_pretrained(self.model_str, src_lang=src_lang)
except:
return M2M100Tokenizer.from_pretrained(self.model_str)
def detect_language(self, text:str)-> str:
"""USes the Lingua package to detect the language of the text.
Args:
text (str): text to be analyzed.
Returns:
str: iso-639-1 code of the detected language.
"""
lang = self.detector.detect_language_of(text)
return lang.iso_code_639_1.name.lower() |