diff --git a/.gitignore b/.gitignore index 70dbcf61..461bc747 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,4 @@ __pycache__/ .coverage .coverage.* *.pyc +.vscode/ diff --git a/src/python/txtai/pipeline/text/translation.py b/src/python/txtai/pipeline/text/translation.py index ade0d126..ad7109fc 100644 --- a/src/python/txtai/pipeline/text/translation.py +++ b/src/python/txtai/pipeline/text/translation.py @@ -114,17 +114,21 @@ class Translation(HFModel): ids = [x.modelId for x in HfApi().list_models(author="Helsinki-NLP")] if self.findmodels else [] return set(ids) - def detect(self, texts): + def detect(self, texts, custom_detect=None): """ Detects the language for each element in texts. Args: texts: list of text + custom_detect: A function to use custom model to detect language Returns: list of languages """ + if custom_detect is not None: + return custom_detect(texts) + if not FASTTEXT: raise ImportError('Language detection is not available - install "pipeline" extra to enable') diff --git a/test/python/testpipeline/testtranslation.py b/test/python/testpipeline/testtranslation.py index ef8a0982..eec3b163 100644 --- a/test/python/testpipeline/testtranslation.py +++ b/test/python/testpipeline/testtranslation.py @@ -13,6 +13,31 @@ class TestTranslation(unittest.TestCase): Translation tests. """ + def testDetect(self): + """ + Test language detection + """ + translate = Translation() + + test = ["This is a test language detection."] + language = translate.detect(test) + + self.assertListEqual(language, ["en"]) + + def testDetectWithCustomFunc(self): + """ + Test language detection with custom function + """ + translate = Translation() + + def dummy_func(text): + return ["en" for x in text] + + test = ["This is a test language detection."] + language = translate.detect(test, dummy_func) + + self.assertListEqual(language, ["en"]) + def testLongTranslation(self): """ Test a translation longer than max tokenization length