diff --git a/examples/custom_models/google_translate_model.py b/examples/custom_models/google_translate_model.py
index 158ee8418..1d67b04e1 100644
--- a/examples/custom_models/google_translate_model.py
+++ b/examples/custom_models/google_translate_model.py
@@ -127,7 +127,7 @@ def greedy_until(
                 # Format is like "community|sdst-text_level:de-fr|0"
                 src_lang, tgt_lang = r.task_name.split("|")[1].split(":")[-1].split("-")
 
-                context = r.context.replace("French phrase: ", "")
+                context = r.context.replace(f"{src_lang.upper()}: ", "").replace(f"\n{tgt_lang.upper()}: ", "")
                 result = self._translate_with_cache(context, src_lang, tgt_lang)
 
                 cur_response = GenerativeResponse(