mirror of
https://github.com/neuml/txtai.git
synced 2026-10-11 21:39:55 +00:00
Add custom detection function for language detection
This commit is contained in:
@@ -7,3 +7,4 @@ __pycache__/
|
|||||||
.coverage
|
.coverage
|
||||||
.coverage.*
|
.coverage.*
|
||||||
*.pyc
|
*.pyc
|
||||||
|
.vscode/
|
||||||
|
|||||||
@@ -114,17 +114,21 @@ class Translation(HFModel):
|
|||||||
ids = [x.modelId for x in HfApi().list_models(author="Helsinki-NLP")] if self.findmodels else []
|
ids = [x.modelId for x in HfApi().list_models(author="Helsinki-NLP")] if self.findmodels else []
|
||||||
return set(ids)
|
return set(ids)
|
||||||
|
|
||||||
def detect(self, texts):
|
def detect(self, texts, custom_detect=None):
|
||||||
"""
|
"""
|
||||||
Detects the language for each element in texts.
|
Detects the language for each element in texts.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
texts: list of text
|
texts: list of text
|
||||||
|
custom_detect: A function to use custom model to detect language
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
list of languages
|
list of languages
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
if custom_detect is not None:
|
||||||
|
return custom_detect(texts)
|
||||||
|
|
||||||
if not FASTTEXT:
|
if not FASTTEXT:
|
||||||
raise ImportError('Language detection is not available - install "pipeline" extra to enable')
|
raise ImportError('Language detection is not available - install "pipeline" extra to enable')
|
||||||
|
|
||||||
|
|||||||
@@ -13,6 +13,31 @@ class TestTranslation(unittest.TestCase):
|
|||||||
Translation tests.
|
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):
|
def testLongTranslation(self):
|
||||||
"""
|
"""
|
||||||
Test a translation longer than max tokenization length
|
Test a translation longer than max tokenization length
|
||||||
|
|||||||
Reference in New Issue
Block a user