complicated lock magic

GooeyAI · Sep 20, 2023 · af6048b · af6048b
1 parent e86c01b
commit af6048b
Show file tree

Hide file tree

Showing 15 changed files with 256 additions and 110 deletions.
diff --git a/daras_ai_v2/asr.py b/daras_ai_v2/asr.py
@@ -17,6 +17,11 @@
     GpuEndpoints,
     call_celery_task,
 )
+from daras_ai_v2.glossary import (
+    DEFAULT_GLOSSARY_URL,
+    LOCATION,
+    glossary_resource,
+)
 
 SHORT_FILE_CUTOFF = 5 * 1024 * 1024  # 1 MB
 
@@ -117,7 +122,7 @@ def google_translate_languages() -> dict[str, str]:
     parent = f"projects/{project}/locations/global"
     client = translate.TranslationServiceClient()
     supported_languages = client.get_supported_languages(
-        parent, display_language_code="en"
+        parent=parent, display_language_code="en"
     )
     return {
         lang.language_code: lang.display_name
@@ -162,6 +167,7 @@ def run_google_translate(
     texts: list[str],
     target_language: str,
     source_language: str = None,
+    glossary_url: str = DEFAULT_GLOSSARY_URL,
 ) -> list[str]:
     """
     Translate text using the Google Translate API.
@@ -189,42 +195,56 @@ def run_google_translate(
     )
 
 
-def _translate_text(text: str, source_language: str, target_language: str):
+def _translate_text(
+    text: str,
+    source_language: str,
+    target_language: str,
+    glossary_url: str = DEFAULT_GLOSSARY_URL,
+) -> str:
     is_romanized = source_language.endswith("-Latn")
     source_language = source_language.replace("-Latn", "")
     enable_transliteration = (
         is_romanized and source_language in TRANSLITERATION_SUPPORTED
     )
+    glossary_url = (
+        glossary_url if not enable_transliteration else ""
+    )  # glossary does not work with transliteration
+
     # prevent incorrect API calls
     if source_language == target_language or not text:
         return text
 
     if source_language == "wo-SN" or target_language == "wo-SN":
         return _MinT_translate_one_text(text, source_language, target_language)
 
-    authed_session, project = get_google_auth_session()
-    res = authed_session.post(
-        f"https://translation.googleapis.com/v3/projects/{project}/locations/global:translateText",
-        json.dumps(
+    config = {
+        "source_language_code": source_language,
+        "target_language_code": target_language,
+        "contents": text,
+        "mime_type": "text/plain",
+        "transliteration_config": {"enable_transliteration": enable_transliteration},
+    }
+
+    with glossary_resource(glossary_url) as (uri, _):
+        config.update(
             {
-                "source_language_code": source_language,
-                "target_language_code": target_language,
-                "contents": text,
-                "mime_type": "text/plain",
-                "transliteration_config": {
-                    "enable_transliteration": enable_transliteration
-                },
+                "glossaryConfig": {
+                    "glossary": uri,
+                    "ignoreCase": True,
+                }
             }
-        ),
-        headers={
-            "Content-Type": "application/json",
-        },
-    )
-    res.raise_for_status()
-    data = res.json()
-    result = data["translations"][0]
+        )
+
+        authed_session, project = get_google_auth_session()
+        res = authed_session.post(
+            f"https://translation.googleapis.com/v3/projects/{project}/locations/{'global' if not uri else LOCATION}:translateText",
+            json=config,
+        )
+        res.raise_for_status()
+        data = res.json()
+        result = data["glossaryTranslations"][0] if uri else data["translations"][0]
 
-    return result["translatedText"].strip()
+        return result["translatedText"].strip()
 
 
 _session = None

diff --git a/daras_ai_v2/doc_search_settings_widgets.py b/daras_ai_v2/doc_search_settings_widgets.py
@@ -17,12 +17,12 @@ def document_uploader(
     accept_multiple_files=True,
 ):
     st.write(label, className="gui-input")
-    documents = st.session_state.get(key) or ([] if accept_multiple_files else "")
-    has_custom_urls = not all(
-        map(is_user_uploaded_url, documents if accept_multiple_files else [documents])
-    )
+    documents = st.session_state.get(key) or []
+    has_custom_urls = not all(map(is_user_uploaded_url, documents))
     custom_key = "__custom_" + key
-    if st.checkbox("Enter Custom URLs", value=has_custom_urls):
+    checkbox_key = "__checkbox_" + key
+    st.session_state.setdefault(checkbox_key, has_custom_urls)
+    if st.checkbox("Enter Custom URLs", key=checkbox_key):
         if not custom_key in st.session_state:
             st.session_state[custom_key] = "\n".join(documents)
         if accept_multiple_files:
@@ -39,7 +39,6 @@ def document_uploader(
                     "fontSize": "0.9rem",
                 },
             )
-            st.session_state[key] = text_value.strip().splitlines()
         else:
             text_value = st.text_input(
                 label,
@@ -53,9 +52,7 @@ def document_uploader(
                     "fontSize": "0.9rem",
                 },
             )
-            st.session_state[key] = (
-                text_value.splitlines()[0] if len(text_value.splitlines()) > 0 else None
-            )
+        st.session_state[key] = text_value.strip().splitlines()
     else:
         st.session_state.pop(custom_key, None)
         st.file_uploader(