add json type option, uspport for gemini

ProbablyFaiz · Jun 24, 2024 · d0f14f7 · d0f14f7
1 parent 917f5f3
commit d0f14f7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/rl/llm/config.py b/rl/llm/config.py
@@ -20,6 +20,7 @@ class LLMConfig:
     frequency_penalty: float = 0.2  # Experiment with this
     num_gpus: int | None = None
     visible_devices: str | None = None
+    json_output: bool = False
 
     def __post_init__(self):
         if not self.tokenizer_name_or_path:

diff --git a/rl/llm/engines.py b/rl/llm/engines.py
@@ -270,7 +270,7 @@ def generate(self, prompt: ChatInput) -> InferenceOutput:
             generation_config={
                 "temperature": self.llm_config.temperature,
                 "max_output_tokens": self.llm_config.max_new_tokens,
-                "response_mime_type": "text/plain",
+                "response_mime_type": "application/json" if self.config.json_output else "text/plain",
             },
             system_instruction=system_message,
             safety_settings={