Register missing engines

ProbablyFaiz · Jun 26, 2024 · bf473de · bf473de
1 parent 7a9e65d
commit bf473de
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/rl/llm/engines.py b/rl/llm/engines.py
@@ -270,7 +270,9 @@ 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": "application/json" if self.llm_config.json_output else "text/plain",
+                "response_mime_type": "application/json"
+                if self.llm_config.json_output
+                else "text/plain",
             },
             system_instruction=system_message,
             safety_settings={
@@ -825,13 +827,15 @@ def _wrap_output(self, req_output) -> InferenceOutput:
     e.NAME: e
     for e in (
         VLLMEngine,
+        AsyncVLLMEngine,
         WorkerVLLMEngine,
         OpenAIEngine,
         TogetherEngine,
         GroqEngine,
         AnthropicEngine,
         ModalEngine,
         GeminiEngine,
+        ManualEditEngine,
     )
 }