Merge pull request #56 from monarch-initiative/deepl-model

Adding DeepL integration to babelon toolkit

pyproject.toml

@@ -11,17 +11,18 @@ python = "^3.11"
 click = "*"
 importlib-metadata = ">=4.8.0"
 rdflib = ">=6.3.2"
-jsonasobj2 = "^1.0.4"
-xmltodict = "^0.13.0"
-pandas = "^2.0.2"
-oaklib = "^0.6.0"
-tabulate = "^0.9.0"
-llm = "^0.13.1"
+jsonasobj2 = ">=1.0.4"
+xmltodict = ">=0.13.0"
+pandas = ">=2.0.2"
+deepl = ">=1.20.0"
+oaklib = ">=0.6.0"
+tabulate = ">=0.9.0"
+llm = ">=0.19.1"
 python-dotenv = "^1.0.1"
-curies = "^0.7.7"
-sssom = "^0.4.4"
-linkml-runtime = "^1.7.1"
-linkml = "^1.7.4"
+curies = ">=0.7.7"
+sssom = ">=0.4.4"
+linkml-runtime = ">=1.7.1"
+linkml = ">=1.7.4"
 
 [tool.poetry.group.docs.dependencies]
 mkdocs = "^1.4.2"

src/babelon/translate.py

@@ -6,6 +6,7 @@
 import string
 from typing import Dict, List
 
+import deepl
 import llm
 import pandas as pd
 
@@ -38,7 +39,7 @@ def translate(self, text, target_language):
 class OpenAITranslator(Translator):
     """A specific translator class that uses GPT-4 for translation."""
 
-    def __init__(self, model="gpt-4-turbo-preview"):
+    def __init__(self, model="gpt-4o"):
         """Instantiate GPT4 translator."""
         self.model = llm.get_model(model)
         self.model.key = os.environ["OPENAI_API_KEY"]
@@ -82,6 +83,40 @@ def translate(self, text_to_translate, language_code):
             return ""
 
 
+class DeepLTranslator(Translator):
+    """A specific translator class that uses DeepL API for translation."""
+
+    def __init__(self):
+        """Instantiate DeepL translator with an API key."""
+        self.api_key = os.environ["DEEPL_API_KEY"]
+        self.translator = deepl.Translator(self.api_key)
+
+    def model_name(self):
+        """Return the unique name of the translation model."""
+        return "DeepL"
+
+    def translate(self, text_to_translate, language_code):
+        """
+        Translate text using DeepL API.
+
+        Args:
+        text_to_translate (str): The text to be translated.
+        language_code (str): The target language code (e.g., 'DE' for German).
+
+        Returns:
+        str: The translated text, or an empty string if translation fails.
+        """
+        result = self.translator.translate_text(
+            text_to_translate, target_lang=language_code.upper()
+        )
+        translation = result.text
+        if translation:
+            print(f"Translation: {translation}")
+            return translation
+        else:
+            return ""
+
+
 def _get_translation_language(translation_language_df, default_language="en"):
     if translation_language_df:
         return translation_language_df
@@ -106,9 +141,11 @@ def get_translator_model(model="gpt-4"):
         ValueError: If the model does not exist.
     """
     if model == "gpt-4":
-        return OpenAITranslator("gpt-4-turbo-preview")
+        return OpenAITranslator("gpt-4o")
     elif model == "gpt-3.5":
         return OpenAITranslator("gpt-3.5-turbo")
+    elif model == "deepl":
+        return DeepLTranslator()
     else:
         try:
             translator = OpenAITranslator(model)

tests/test_translate.py

@@ -7,6 +7,7 @@
 from oaklib import get_adapter
 
 from babelon.translate import (
+    DeepLTranslator,
     OpenAITranslator,
     _is_equivalent_string,
     prepare_translation_for_ontology,
@@ -24,13 +25,21 @@ def setUp(self) -> None:
         """Set up the test case."""
 
     @unittest.skipIf(not os.path.exists(env_file), "Skipping test as .env file does not exist")
-    def test_translate(self):
+    def test_translate_openai(self):
         """Test update_translation_profile."""
         load_dotenv()
         translator = OpenAITranslator()
         translated_value = translator.translate("fever", "de")
         self.assertEqual("Fieber", translated_value)
 
+    @unittest.skipIf(not os.path.exists(env_file), "Skipping test as .env file does not exist")
+    def test_translate_deepl(self):
+        """Test update_translation_profile."""
+        load_dotenv()
+        translator = DeepLTranslator()
+        translated_value = translator.translate("fever", "de")
+        self.assertEqual("Fieber", translated_value)
+
     @unittest.skipIf(not os.path.exists(env_file), "Skipping test as .env file does not exist")
     def test_translate_profile(self):
         """Test to see if a small babelon profile can be translated."""