Fixed problem with removing languages in the context.

huggingface · Dec 12, 2024 · d1af518 · d1af518
1 parent 71f871e
commit d1af518
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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(