better handle missing src and domain tokens

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers.cache_utils import DynamicCache
 import torch
 import itertools
 
@@ -75,11 +76,11 @@ CODE2LANG = {v: k for k, v in LANG2CODE.items()}
 LANGUAGES = sorted(LANG2CODE.keys())
 
 
-def language_token(lang):
+def build_language_token(lang):
     return f"<lang_{lang}>"
 
 
-def domain_token(dom):
+def build_domain_token(dom):
     return f"<dom_{dom}>"
 
 
@@ -92,7 +93,7 @@ def domain_token_to_str(token):
 
 
 def format_input(src, tgt_lang, src_lang, domain):
-    tgt_lang_token = language_token(tgt_lang)
+    tgt_lang_token = build_language_token(tgt_lang)
 
     prefix = TOKENIZER.eos_token
 
@@ -100,13 +101,13 @@ def format_input(src, tgt_lang, src_lang, domain):
     if src_lang is None:
         return base_input
     else:
-        src_lang_token = language_token(src_lang)
+        src_lang_token = build_language_token(src_lang)
         base_input = f"{base_input}{src_lang_token}"
 
     if domain is None:
         return base_input
     else:
-        dom_token = domain_token(domain)
+        dom_token = build_domain_token(domain)
         base_input = f"{base_input}{dom_token}"
 
     return base_input
@@ -115,27 +116,109 @@ def format_input(src, tgt_lang, src_lang, domain):
 def translate_with_model(model_name, text, tgt_lang, src_lang, domain):
     model = MODELS[model_name]
     formatted_text = format_input(text, tgt_lang, src_lang, domain)
-
-    inputs = TOKENIZER(formatted_text, return_tensors="pt", return_token_type_ids=False)
+    inputs = TOKENIZER(
+        formatted_text,
+        return_attention_mask=True,
+        return_tensors="pt",
+        return_token_type_ids=False,
+    )
     for k, v in inputs.items():
         inputs[k] = v.to(DEVICE)
+    src_lang_provided = src_lang is not None
+    domain_provided = domain is not None
+    need_format_again = not (src_lang_provided and domain_provided)
 
-    if src_lang is None:
-        domain_token_pos = inputs["input_ids"].size(1) + 1
-    elif domain is None:
-        domain_token_pos = inputs["input_ids"].size(1)
-    else:
-        domain_token_pos = inputs["input_ids"].size(1) - 1
+    past_key_values = DynamicCache()
+    cache_position = torch.arange(
+        inputs["input_ids"].size(1), dtype=torch.int64, device=DEVICE
+    )
+
+    if not src_lang_provided:
+        # Need to predict src lang
+        with torch.inference_mode():
+            outputs = model(
+                input_ids=inputs["input_ids"],
+                attention_mask=inputs["attention_mask"],
+                use_cache=True,
+                past_key_values=past_key_values,
+                cache_position=cache_position,
+            )
+            src_lang_token_id = outputs.logits[:, -1, :].argmax(dim=-1).unsqueeze(0)
+            src_lang = language_token_to_str(
+                TOKENIZER.convert_ids_to_tokens(src_lang_token_id.squeeze().item())
+            )
+
+            cache_position = cache_position[-1:] + 1
+
+            attention_mask = inputs["attention_mask"]
+            attention_mask = torch.cat(
+                [attention_mask, attention_mask.new_ones((attention_mask.size(0), 1))],
+                dim=-1,
+            )
+            inputs = {"input_ids": src_lang_token_id, "attention_mask": attention_mask}
+
+    if not domain_provided:
+        # Need to predict domain
+        with torch.inference_mode():
+            outputs = model(
+                input_ids=inputs["input_ids"],
+                attention_mask=inputs["attention_mask"],
+                use_cache=True,
+                past_key_values=past_key_values,
+            )
+            domain_token_id = outputs.logits[:, -1, :].argmax(dim=-1).unsqueeze(0)
+            domain = domain_token_to_str(
+                TOKENIZER.convert_ids_to_tokens(domain_token_id.squeeze().item())
+            )
+
+            cache_position = cache_position[-1:] + 1
+
+            attention_mask = inputs["attention_mask"]
+            attention_mask = torch.cat(
+                [attention_mask, attention_mask.new_ones((attention_mask.size(0), 1))],
+                dim=-1,
+            )
+
+            inputs = {"input_ids": domain_token_id, "attention_mask": attention_mask}
+    elif not src_lang_provided:
+        # in this case, src_lang was not provided, but domain was.
+        # So we still need to run a forward pass to build the kv cache for the domain token
+        dom_token = build_domain_token(domain)
+        # dom_token = "<dom_general>"
+        domain = domain_token_to_str(dom_token)
+
+        domain_token_id = TOKENIZER.convert_tokens_to_ids(dom_token)
+        inputs["input_ids"] = torch.hstack(
+            [inputs["input_ids"], torch.tensor([[domain_token_id]], device=DEVICE)]
+        )
+        inputs["attention_mask"] = torch.hstack(
+            [inputs["attention_mask"], inputs["attention_mask"].new_ones((1, 1))]
+        )
+        cache_position = torch.hstack([cache_position, cache_position[-1:] + 1])
+
+    if need_format_again:
+        formatted_text = format_input(text, tgt_lang, src_lang, domain)
+        inputs = TOKENIZER(
+            formatted_text,
+            return_attention_mask=True,
+            return_tensors="pt",
+            return_token_type_ids=False,
+        )
+        for k, v in inputs.items():
+            inputs[k] = v.to(DEVICE)
+
+    domain_token_pos = inputs["input_ids"].size(1) - 1
     src_lang_token_pos = domain_token_pos - 1
     _tgt_lang_token_pos = src_lang_token_pos - 1
 
     outputs = model.generate(
-        **inputs,
-        num_beams=5,
-        length_penalty=0.65,
+        input_ids=inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
+        num_beams=1,
         max_new_tokens=500,
         pad_token_id=TOKENIZER.pad_token_id,
         eos_token_id=TOKENIZER.eos_token_id,
+        past_key_values=past_key_values,
     )
 
     generated_translation = TOKENIZER.decode(
@@ -145,12 +228,12 @@ def translate_with_model(model_name, text, tgt_lang, src_lang, domain):
     source_language_token = TOKENIZER.convert_ids_to_tokens(
         outputs[0, src_lang_token_pos].item()
     )
-    domain_token = TOKENIZER.convert_ids_to_tokens(outputs[0, domain_token_pos].item())
+    dom_token = TOKENIZER.convert_ids_to_tokens(outputs[0, domain_token_pos].item())
 
     return {
         "model": model_name,
         "source_lang": CODE2LANG[language_token_to_str(source_language_token)],
-        "domain": DOMAIN_MAPPING_REVERSED[domain_token_to_str(domain_token)],
+        "domain": DOMAIN_MAPPING_REVERSED[domain_token_to_str(dom_token)],
         "translation": generated_translation,
     }