Fix role

ProbablyFaiz · Jun 19, 2024 · de4fcf0 · de4fcf0
1 parent c9d32a4
commit de4fcf0
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/rl/llm/engines.py b/rl/llm/engines.py
@@ -304,7 +304,13 @@ def _convert_openai_to_gemini(
         prompt = prompt[:-1]
         return (
             system_prompt,
-            [{"role": msg["role"], "parts": [msg["content"]]} for msg in prompt],
+            [
+                {
+                    "role": "model" if msg["role"] == "assistant" else msg["role"],
+                    "parts": [msg["content"]],
+                }
+                for msg in prompt
+            ],
             last_message,
         )