119 lines
3.9 KiB
Python
119 lines
3.9 KiB
Python
"""
|
|
pybergamot - (Somewhat) stable interface for the **Bergamot Translation Engine Python Bindings**.
|
|
Copyright (C) 2023 Ad5001
|
|
|
|
This program is free software: you can redistribute it and/or modify
|
|
it under the terms of the GNU General Public License as published by
|
|
the Free Software Foundation, either version 3 of the License, or
|
|
(at your option) any later version.
|
|
|
|
This program is distributed in the hope that it will be useful,
|
|
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
GNU General Public License for more details.
|
|
|
|
You should have received a copy of the GNU General Public License
|
|
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
|
|
Meta repositories for translation engines
|
|
"""
|
|
from abc import ABC, abstractmethod
|
|
from bergamot import Service, VectorString, ResponseOptions, VectorResponse, TranslationModel
|
|
|
|
|
|
class Engine(ABC):
|
|
"""
|
|
An interface for multiple types of translation engine.
|
|
"""
|
|
|
|
@property
|
|
@abstractmethod
|
|
def source_lang(self) -> str:
|
|
"""
|
|
Two-char ISO languages for which the engine translates from.
|
|
"""
|
|
pass
|
|
|
|
@property
|
|
@abstractmethod
|
|
def target_lang(self) -> str:
|
|
"""
|
|
Two-char ISO languages for which the engine translates to.
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def translate(self, text: str, html: bool = False, alignment: bool = False, quality_scores: bool = False) -> str:
|
|
"""
|
|
Translates the text from the engine's source lang its target.
|
|
|
|
## Parameters
|
|
---
|
|
- **text**: Text to translate.
|
|
- **html**: Set to True if the text contains an HTML structure which needs to
|
|
be preserved while translated.
|
|
- **alignment**: Toggle for alignment.
|
|
- **quality_scores**: Toggle for whether to include the translation's quality scores
|
|
for each word in HTML format.
|
|
|
|
## Returns
|
|
---
|
|
The translated text.
|
|
"""
|
|
pass
|
|
|
|
|
|
class DirectBergamotModelEngine(Engine):
|
|
"""
|
|
Internal module. Engine using a single bergamot model to translate.
|
|
"""
|
|
def __init__(self, source_lang: str, target_lang: str, model: TranslationModel,
|
|
service: Service):
|
|
self._source_lang = source_lang
|
|
self._target_lang = target_lang
|
|
self.model = model
|
|
self.service = service
|
|
|
|
@property
|
|
def source_lang(self) -> str:
|
|
return self._source_lang
|
|
|
|
@property
|
|
def target_lang(self) -> str:
|
|
return self._target_lang
|
|
|
|
def translate(self, text: str, html: bool = False, alignment: bool = False, quality_scores: bool = False) -> str:
|
|
opts = ResponseOptions(
|
|
alignment=alignment, qualityScores=quality_scores, HTML=html
|
|
)
|
|
resp: VectorResponse = self.service.translate(self.model, VectorString([text]), opts)
|
|
return resp[0].target.text
|
|
|
|
|
|
class ChainBergamotModelsEngine(Engine):
|
|
"""
|
|
Internal module. Engine chaining two bergamot model to translate.
|
|
"""
|
|
def __init__(self, source_lang: str, target_lang: str, model1: TranslationModel,
|
|
model2: TranslationModel, service: Service):
|
|
self._source_lang = source_lang
|
|
self._target_lang = target_lang
|
|
self.model1 = model1
|
|
self.model2 = model2
|
|
self.service = service
|
|
|
|
@property
|
|
def source_lang(self) -> str:
|
|
return self._source_lang
|
|
|
|
@property
|
|
def target_lang(self) -> str:
|
|
return self._target_lang
|
|
|
|
def translate(self, text: str, html: bool = False, alignment: bool = False, quality_scores: bool = False) -> str:
|
|
opts = ResponseOptions(
|
|
alignment=alignment, qualityScores=quality_scores, HTML=html
|
|
)
|
|
resp: VectorResponse = self.service.pivot(self.model1, self.model2, VectorString([text]), opts)
|
|
return resp[0].target.text
|