Wire in signwriting-to-text model

Dockerfile

@@ -20,7 +20,10 @@ RUN timeout 360 bash -c '\
     curl --fail --retry 10 --retry-delay 5 --retry-connrefused \
          --location --request POST "http://localhost:8080/" \
          --header "Content-Type: application/json" \
-         --data "{\"texts\": [\"hello\"], \"spoken_language\": \"en\", \"signed_language\": \"ase\"}"; \
+         --data "{\"texts\": [\"hello\"], \"spoken_language\": \"en\", \"signed_language\": \"ase\"}" && \
+    curl --fail --location --request POST "http://localhost:8080/signwriting-to-text" \
+         --header "Content-Type: application/json" \
+         --data "{\"texts\": [\"M528x518S15a28472x487S1f010490x502\"], \"spoken_language\": \"en\", \"signed_language\": \"ase\"}"; \
     CURL_EXIT=$?; \
     kill $SERVER_PID 2>/dev/null || true; \
     wait $SERVER_PID 2>/dev/null || true; \

signwriting_translation/bin.py

@@ -3,9 +3,10 @@
 import argparse
 import time
 from functools import cache
+from itertools import chain
 from pathlib import Path
 
-from signwriting.tokenizer import SignWritingTokenizer
+from signwriting.tokenizer import SignWritingTokenizer, normalize_signwriting
 from sockeye.inference import TranslatorOutput
 
 from signwriting_translation.tokenizer import tokenize_spoken_text
@@ -19,6 +20,16 @@ def process_translation_output(output: TranslatorOutput):
     return sw_tokenizer.tokens_to_text((" ".join(symbols)).split(" "))
 
 
+def tokenize_signwriting(fsw: str) -> str:
+    signs = normalize_signwriting(fsw).split(" ")
+    return " ".join(chain.from_iterable(sw_tokenizer.text_to_tokens(sign) for sign in signs))
+
+
+def process_text_output(output: TranslatorOutput) -> str:
+    # target side is subword-nmt BPE; "@@" marks a split within a word
+    return " ".join(output.tokens).replace("@@ ", "").strip()
+
+
 @cache
 def load_sockeye_translator(model_path: str, log_timing: bool = False):
     if not Path(model_path).is_dir():
@@ -41,7 +52,7 @@ def load_sockeye_translator(model_path: str, log_timing: bool = False):
     return translator, tokenizer_path
 
 
-def translate(translator, texts: list[str], log_timing: bool = False):
+def _run_translator(translator, texts: list[str], log_timing: bool = False):
     from sockeye.inference import make_input_from_plain_string
 
     inputs = [make_input_from_plain_string(sentence_id=i, string=s)
@@ -50,36 +61,48 @@ def translate(translator, texts: list[str], log_timing: bool = False):
     now = time.time()
     outputs = translator.translate(inputs)
     translation_time = time.time() - now
-    avg_time = translation_time / len(texts)
     if log_timing:
+        avg_time = translation_time / len(texts)
         print("Translated", len(texts), "texts in", translation_time, "seconds", f"({avg_time:.2f} seconds per text)")
+    return outputs
+
+
+def translate(translator, texts: list[str], log_timing: bool = False):
+    outputs = _run_translator(translator, texts, log_timing)
     return [process_translation_output(output) for output in outputs]
 
 
-def get_args():
+def translate_to_text(translator, texts: list[str], log_timing: bool = False):
+    outputs = _run_translator(translator, texts, log_timing)
+    return [process_text_output(output) for output in outputs]
+
+
+def get_args(default_model: str):
     parser = argparse.ArgumentParser()
-    parser.add_argument('--model', type=str, help='Path to trained model',
-                        default="sign/sockeye-text-to-factored-signwriting")
+    parser.add_argument('--model', type=str, help='Path to trained model', default=default_model)
     parser.add_argument('--spoken-language', required=True, type=str, help='spoken language code')
     parser.add_argument('--signed-language', required=True, type=str, help='signed language code')
     parser.add_argument('--input', required=True, type=str, help='input text or signwriting sequence')
     return parser.parse_args()
 
 
 def text_to_signwriting():
-    args = get_args()
+    args = get_args("sign/sockeye-text-to-factored-signwriting")
 
-    translator, tokenizer_path = load_sockeye_translator(args.model)
+    translator, _ = load_sockeye_translator(args.model)
     tokenized_text = tokenize_spoken_text(args.input)
     model_input = f"${args.spoken_language} ${args.signed_language} {tokenized_text}"
     outputs = translate(translator, [model_input])
     print(outputs[0])
 
 
 def signwriting_to_text():
-    args = get_args()
+    args = get_args("sign/sockeye-signwriting-to-text")
 
-    return translate(args.model, [args.input])
+    translator, _ = load_sockeye_translator(args.model)
+    model_input = f"${args.spoken_language} ${args.signed_language} {tokenize_signwriting(args.input)}"
+    outputs = translate_to_text(translator, [model_input])
+    print(outputs[0])
 
 
 if __name__ == '__main__':

signwriting_translation/server.py

@@ -5,12 +5,19 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 
-from signwriting_translation.bin import load_sockeye_translator, translate
+from signwriting_translation.bin import (
+    load_sockeye_translator,
+    tokenize_signwriting,
+    translate,
+    translate_to_text,
+)
 from signwriting_translation.tokenizer import tokenize_spoken_text
 
-MODEL_ID = "sign/sockeye-text-to-factored-signwriting"
+TEXT_TO_SW_MODEL = "sign/sockeye-text-to-factored-signwriting"
+SW_TO_TEXT_MODEL = "sign/sockeye-signwriting-to-text"
 
-translator, _ = load_sockeye_translator(MODEL_ID, log_timing=True)
+text_to_sw_translator, _ = load_sockeye_translator(TEXT_TO_SW_MODEL, log_timing=True)
+sw_to_text_translator, _ = load_sockeye_translator(SW_TO_TEXT_MODEL, log_timing=True)
 
 app = FastAPI(title="Signwriting Translation API")
 
@@ -40,7 +47,22 @@ def spoken_text_to_signwriting(request: TranslationRequest):
         for text in request.texts
     ]
 
-    output_texts = translate(translator, tokenized_texts, log_timing=True)
+    output_texts = translate(text_to_sw_translator, tokenized_texts, log_timing=True)
+
+    return {"input": request.texts, "output": output_texts}
+
+
+@app.post("/signwriting-to-text")
+def signwriting_to_spoken_text(request: TranslationRequest):
+    if not request.texts:
+        raise HTTPException(status_code=400, detail="Missing `texts`")
+
+    tokenized_texts = [
+        f"${request.spoken_language} ${request.signed_language} {tokenize_signwriting(text)}"
+        for text in request.texts
+    ]
+
+    output_texts = translate_to_text(sw_to_text_translator, tokenized_texts, log_timing=True)
 
     return {"input": request.texts, "output": output_texts}
 

tests/test_server.py

@@ -27,6 +27,19 @@ def test_translate():
     assert body["output"][0].startswith("M")
 
 
+def test_translate_signwriting_to_text():
+    response = client.post("/signwriting-to-text", json={
+        "texts": ["M528x518S15a28472x487S1f010490x502"],
+        "spoken_language": "en",
+        "signed_language": "ase",
+    })
+    assert response.status_code == 200
+    body = response.json()
+    assert len(body["output"]) == 1
+    assert isinstance(body["output"][0], str)
+    assert body["output"][0]
+
+
 def test_translate_empty_texts():
     response = client.post("/", json={
         "texts": [],