add multilingual support and cache translations

app.py

@@ -8,24 +8,26 @@ from transformers import LogitsProcessorList
 
 import streamlit as st
 
-from utils import get_resources
+from utils import get_resources, LANGUAGES, translate
 
-model, tokenizer, logitsprocessor = get_resources()
+st.title("👩‍💻 Generate AMR from multilingual text")
 
-st.title("📝 Parse text into AMR")
+with st.form("input data"):
+    text_col, lang_col = st.columns((4, 1))
+    text = text_col.text_input(label="Input text")
+    src_lang = lang_col.selectbox(label="Language", options=list(LANGUAGES.keys()), index=0)
+    submitted = st.form_submit_button("Submit")
 
-text = st.text_input(label="Text to transform (en)")
-
-if text:
+if submitted:
+    multilingual = src_lang != "English"
+    model, tokenizer, logitsprocessor = get_resources(multilingual)
     gen_kwargs = {
         "max_length": model.config.max_length,
         "num_beams": model.config.num_beams,
         "logits_processor": LogitsProcessorList([logitsprocessor])
     }
 
-    encoded = tokenizer(text, return_tensors="pt")
-    generated = model.generate(**encoded, **gen_kwargs)
-    linearized = tokenizer.decode_and_fix(generated)[0]
+    linearized = translate(text, src_lang, model, tokenizer, **gen_kwargs)
     penman_str = linearized2penmanstr(linearized)
 
     try:

utils.py

@@ -16,9 +16,15 @@ st_hash_funcs = {PreTrainedModel: lambda model: model.name_or_path,
 
 
 @st.cache(show_spinner=False, hash_funcs=st_hash_funcs, allow_output_mutation=True)
-def get_resources(quantize: bool = True):
-    tokenizer = AMRMBartTokenizer.from_pretrained("BramVanroy/mbart-en-to-amr", src_lang="en_XX")
-    model = MBartForConditionalGeneration.from_pretrained("BramVanroy/mbart-en-to-amr")
+def get_resources(multilingual: bool, quantize: bool = True):
+    if multilingual:
+        # Tokenizer src_lang is reset during translation to the right language
+        tokenizer = AMRMBartTokenizer.from_pretrained("BramVanroy/mbart-en-es-nl-to-amr", src_lang="nl_XX")
+        model = MBartForConditionalGeneration.from_pretrained("BramVanroy/mbart-en-es-nl-to-amr")
+    else:
+        tokenizer = AMRMBartTokenizer.from_pretrained("BramVanroy/mbart-en-to-amr", src_lang="en_XX")
+        model = MBartForConditionalGeneration.from_pretrained("BramVanroy/mbart-en-to-amr")
+
     model = BetterTransformer.transform(model, keep_original_model=False)
     model.resize_token_embeddings(len(tokenizer))
 
@@ -28,3 +34,20 @@ def get_resources(quantize: bool = True):
     logits_processor = AMRLogitsProcessor(tokenizer, model.config.max_length)
 
     return model, tokenizer, logits_processor
+
+
+@st.cache(show_spinner=False, hash_funcs=st_hash_funcs)
+def translate(text: str, src_lang: str, model: MBartForConditionalGeneration, tokenizer: AMRMBartTokenizer, **gen_kwargs):
+    tokenizer.src_lang = LANGUAGES[src_lang]
+    encoded = tokenizer(text, return_tensors="pt")
+    print(tokenizer.convert_ids_to_tokens(encoded.input_ids[0]))
+    print(model.name_or_path)
+    generated = model.generate(**encoded, **gen_kwargs)
+    return tokenizer.decode_and_fix(generated)[0]
+
+
+LANGUAGES = {
+    "English": "en_XX",
+    "Dutch": "nl_XX",
+    "Spanish": "es_XX",
+}