Claude/add ollama logprobs support eu qxg

atroposlib/envs/server_handling/server_baseline.py

@@ -108,7 +108,7 @@ class ServerBaseline(BaseModel):
     rolling_buffer_length: int = Field(
         default=1000, description="Length of the rolling buffer to store metrics."
     )
-    server_type: Literal["openai", "trl", "sglang", "vllm"] = Field(
+    server_type: Literal["openai", "trl", "sglang", "vllm", "ollama"] = Field(
         default="openai", description="Type of server to use"
     )
 

atroposlib/envs/server_handling/server_manager.py

@@ -10,6 +10,7 @@
 from pydantic import BaseModel, Field
 
 from atroposlib.envs.server_handling.managed_server import ManagedServer
+from atroposlib.envs.server_handling.ollama_server import OllamaServer
 from atroposlib.envs.server_handling.openai_server import OpenAIServer
 from atroposlib.envs.server_handling.server_baseline import (
     APIServer,
@@ -62,6 +63,8 @@ def __init__(
                     server_class = SGLangServer
                 elif configs.server_type == "vllm":
                     server_class = VLLMServer
+                elif configs.server_type == "ollama":
+                    server_class = OllamaServer
                 else:
                     raise ValueError(f"Invalid server type: {configs.server_type}")
             else:
@@ -73,6 +76,8 @@ def __init__(
                     server_class = SGLangServer
                 elif configs[0].server_type == "vllm":
                     server_class = VLLMServer
+                elif configs[0].server_type == "ollama":
+                    server_class = OllamaServer
                 else:
                     raise ValueError(f"Invalid server type: {configs[0].server_type}")
         if testing: