Add API endpoint

api_endpoint.py

@@ -0,0 +1,34 @@
+import json
+import os
+import requests
+from traceback import print_exc
+from dotenv import load_dotenv
+
+load_dotenv()
+
+api_endpoint_url = os.environ.get("API_ENDPOINT_URL")
+api_endpoint_basic_auth_user = os.environ.get("API_ENDPOINT_BASIC_AUTH_USER")
+api_endpoint_basic_auth_password = os.environ.get("API_ENDPOINT_BASIC_AUTH_PASSWORD")
+api_endpoint_enabled = os.environ.get("API_ENDPOINT_ENABLED", default=True) == "True"
+
+def invoke_translate_endpoint(text, model):
+    try:
+        response = requests.request("POST", 
+                                    f"{api_endpoint_url}/api/v1/translate",
+                                    headers={
+                                        'Content-Type': 'application/json', 
+                                    },
+                                    auth=(api_endpoint_basic_auth_user, api_endpoint_basic_auth_password),
+                                    data=json.dumps({
+                                        "src": model.get('src'),
+                                        "tgt": model.get('tgt') ,
+                                        "text": text
+                                        })
+                                    )
+        response.raise_for_status()
+
+        return response.json()['translation']
+
+    except:
+        print_exc()
+        return None

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 from AinaTheme import theme
 from huggingface_hub import snapshot_download
 import nltk
+from api_endpoint import invoke_translate_endpoint
 
 from translate import translate_text
 
@@ -13,7 +14,9 @@ load_dotenv()
 
 MODELS_PATH = "./models"
 HF_CACHE_DIR = "./hf_cache"
-MAX_INPUT_CHARACTERS= int(os.environ.get("MAX_INPUT_CHARACTERS", default=1000))
+MAX_INPUT_CHARACTERS = int(os.environ.get("MAX_INPUT_CHARACTERS", default=1000))
+API_ENDPOINT_ENABLED = os.environ.get("API_ENDPOINT_ENABLED", default=True) == "True"
+
 
 def download_model(repo_id, revision="main"):
     return snapshot_download(repo_id=repo_id, revision=revision, local_dir=os.path.join(MODELS_PATH, repo_id), cache_dir=HF_CACHE_DIR)
@@ -37,7 +40,7 @@ model_dir_pt_ca = download_model("projecte-aina/mt-aina-pt-ca", revision="main")
 directions = {
     "Catalan": {
         "target": {
-            "Spanish": {"model": (f"{os.path.join(MODELS_PATH, model_dir_ca_es)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_ca_es)}")},
+            "Spanish": {"src": "ca", "tgt":"es","model": (f"{os.path.join(MODELS_PATH, model_dir_ca_es)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_ca_es)}")},
             "English": {"model": (f"{os.path.join(MODELS_PATH, model_dir_ca_en)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_ca_en)}")},
             "French": {"model": (f"{os.path.join(MODELS_PATH, model_dir_ca_fr)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_ca_fr)}")},
             "German": {"model": (f"{os.path.join(MODELS_PATH, model_dir_ca_de)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_ca_de)}")},
@@ -48,7 +51,7 @@ directions = {
     },
     "Spanish": {
         "target": {
-            "Catalan": {"model": (f"{os.path.join(MODELS_PATH, model_dir_es_ca)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_es_ca)}")},
+            "Catalan": {"src": "es", "tgt":"ca","model": (f"{os.path.join(MODELS_PATH, model_dir_es_ca)}/spm.model", f"{os.path.join(MODELS_PATH, model_dir_es_ca)}")},
         }
     },
     "English": {
@@ -75,7 +78,8 @@ def get_target_languages(source_language):
 
 
 def get_target_languege_model(source_language, target_language):
-    return directions.get(source_language, {}).get("target", {}).get(target_language, {}).get("model")
+    # return directions.get(source_language, {}).get("target", {}).get(target_language, {}).get("model")
+    return directions.get(source_language, {}).get("target", {}).get(target_language, {})
 
 
 def translate(text, source_language, lang_pair):
@@ -101,7 +105,12 @@ def translate_input(input, source_language, target_language):
         return None
     
     target_language_model = get_target_languege_model(source_language, target_language)
-    translation = translate(input, source_language, target_language_model)
+
+    if (API_ENDPOINT_ENABLED and 'src' in target_language_model):
+         translation =  invoke_translate_endpoint(input, target_language_model)
+         if translation is not None:
+            return translation
+    translation = translate(input, source_language, target_language_model.get('model'))
 
     return translation