Make one model selectable by user (#23)

app.py

@@ -363,6 +363,12 @@ Every benchmark is limited in some sense -- Before you interpret the results, pl
 controller_addr = os.environ["COLOSSEUM_CONTROLLER_ADDR"]
 global_controller_client = ControllerClient(controller_addr=controller_addr, timeout=15)
 
+# Load the list of models. To reload, the app should be restarted.
+available_models = global_controller_client.get_available_models()
+model_preference_dropdown_choices = [f"One is {model}" for model in available_models]
+model_preference_dropdown_choices = ["Two random models"] + model_preference_dropdown_choices
+user_pref_to_model_name = dict(zip(model_preference_dropdown_choices, ["Random"] + available_models))
+
 # Colosseum helper functions.
 def enable_interact():
     return [gr.update(interactive=True)] * 2
@@ -394,21 +400,23 @@ def consumed_more_energy_message(energy_a, energy_b):
 
 # Colosseum event handlers
 def add_prompt_disable_submit(prompt, history_a, history_b):
-    """Add the user's prompt to the two model's history and disable the submit button."""
+    """Add the user's prompt to the two model's history and disable further submission."""
     client = global_controller_client.fork()
     return [
         gr.Textbox.update(value=" ", interactive=False),
         gr.Button.update(interactive=False),
+        gr.Dropdown.update(interactive=False),
         history_a + [[prompt, ""]],
         history_b + [[prompt, ""]],
         client,
     ]
 
-def generate_responses(client: ControllerClient, history_a, history_b):
+def generate_responses(client: ControllerClient, user_preference, history_a, history_b):
     """Generate responses for the two models."""
+    model_preference = user_pref_to_model_name[user_preference]
     for resp_a, resp_b in itertools.zip_longest(
-        client.prompt(prompt=history_a[-1][0], index=0),
-        client.prompt(prompt=history_b[-1][0], index=1),
+        client.prompt(prompt=history_a[-1][0], index=0, model_preference=model_preference),
+        client.prompt(prompt=history_b[-1][0], index=1, model_preference=model_preference),
     ):
         if resp_a is not None:
             history_a[-1][1] += resp_a
@@ -475,12 +483,14 @@ def play_again():
     return [
         # Clear chatbot history
         None, None,
-        # Turn on prompt textbox and submit button
+        # Enable prompt textbox and submit button
         gr.Textbox.update(value="", interactive=True), gr.Button.update(interactive=True),
         # Mask model names
         gr.Markdown.update(value="", visible=False), gr.Markdown.update(value="", visible=False),
         # Hide energy vote buttons and message
         gr.Button.update(visible=False), gr.Button.update(visible=False), gr.Markdown.update(visible=False),
+        # Enable model preference dropdown
+        gr.Dropdown.update(interactive=True),
         # Disable reset button
         gr.Button.update(interactive=False, visible=False),
     ]
@@ -506,6 +516,14 @@ with gr.Blocks(css=custom_css) as block:
         with gr.TabItem("Colosseum ⚔️️"):
             gr.Markdown(open("docs/colosseum_top.md").read())
 
+            with gr.Row():
+                model_preference_dropdown = gr.Dropdown(
+                    choices=model_preference_dropdown_choices,
+                    value=model_preference_dropdown_choices[0],
+                    label="Prefer a specific model?",
+                    interactive=True,
+                )
+
             with gr.Group():
                 with gr.Row():
                     prompt_input = gr.Textbox(
@@ -561,12 +579,12 @@ with gr.Blocks(css=custom_css) as block:
 
 
             (prompt_input
-                .submit(add_prompt_disable_submit, [prompt_input, *chatbots], [prompt_input, prompt_submit_btn, *chatbots, controller_client], queue=False)
-                .then(generate_responses, [controller_client, *chatbots], [*chatbots], queue=True, show_progress="hidden")
+                .submit(add_prompt_disable_submit, [prompt_input, *chatbots], [prompt_input, prompt_submit_btn, model_preference_dropdown, *chatbots, controller_client], queue=False)
+                .then(generate_responses, [controller_client, model_preference_dropdown, *chatbots], [*chatbots], queue=True, show_progress="hidden")
                 .then(enable_interact, None, resp_vote_btn_list, queue=False))
             (prompt_submit_btn
-                .click(add_prompt_disable_submit, [prompt_input, *chatbots], [prompt_input, prompt_submit_btn, *chatbots, controller_client], queue=False)
-                .then(generate_responses, [controller_client, *chatbots], [*chatbots], queue=True, show_progress="hidden")
+                .click(add_prompt_disable_submit, [prompt_input, *chatbots], [prompt_input, prompt_submit_btn, model_preference_dropdown, *chatbots, controller_client], queue=False)
+                .then(generate_responses, [controller_client, model_preference_dropdown, *chatbots], [*chatbots], queue=True, show_progress="hidden")
                 .then(enable_interact, None, resp_vote_btn_list, queue=False))
 
             left_resp_vote_btn.click(
@@ -599,7 +617,7 @@ with gr.Blocks(css=custom_css) as block:
                 .click(
                     play_again,
                     None,
-                    [*chatbots, prompt_input, prompt_submit_btn, *masked_model_names, *energy_vote_btn_list, energy_comparison_message, play_again_btn],
+                    [*chatbots, prompt_input, prompt_submit_btn, *masked_model_names, *energy_vote_btn_list, energy_comparison_message, model_preference_dropdown, play_again_btn],
                     queue=False,
                 )
                 .then(None, _js=focus_prompt_input_js, queue=False))

spitfight/colosseum/client.py

@@ -9,9 +9,11 @@ import requests
 import gradio as gr
 
 from spitfight.colosseum.common import (
+    COLOSSEUM_MODELS_ROUTE,
     COLOSSEUM_PROMPT_ROUTE,
     COLOSSEUM_RESP_VOTE_ROUTE,
     COLOSSEUM_ENERGY_VOTE_ROUTE,
+    ModelsResponse,
     PromptRequest,
     ResponseVoteRequest,
     ResponseVoteResponse,
@@ -37,9 +39,33 @@ class ControllerClient:
             request_id=uuid4(),
         )
 
-    def prompt(self, prompt: str, index: Literal[0, 1]) -> Generator[str, None, None]:
-        """Generate the response of the `index`th model with the prompt."""
-        prompt_request = PromptRequest(request_id=self.request_id, prompt=prompt, model_index=index)
+    def get_available_models(self) -> list[str]:
+        """Retrieve the list of available models."""
+        with _catch_requests_exceptions():
+            resp = requests.get(
+                f"http://{self.controller_addr}{COLOSSEUM_MODELS_ROUTE}",
+                timeout=self.timeout,
+            )
+        _check_response(resp)
+        return ModelsResponse(**resp.json()).available_models
+
+    def prompt(
+        self,
+        prompt: str,
+        index: Literal[0, 1],
+        model_preference: str,
+    ) -> Generator[str, None, None]:
+        """Generate the response of the `index`th model with the prompt.
+
+        `user_pref` is the user's preference for the model to use. It can be
+        `"random"` or one of the models in the list returned by `get_available_models`.
+        """
+        prompt_request = PromptRequest(
+            request_id=self.request_id,
+            prompt=prompt,
+            model_index=index,
+            model_preference=model_preference,
+        )
         with _catch_requests_exceptions():
             resp = requests.post(
                 f"http://{self.controller_addr}{COLOSSEUM_PROMPT_ROUTE}",

spitfight/colosseum/common.py

@@ -4,16 +4,22 @@ from typing import Literal
 
 from pydantic import BaseModel
 
+COLOSSEUM_MODELS_ROUTE = "/models"
 COLOSSEUM_PROMPT_ROUTE = "/prompt"
 COLOSSEUM_RESP_VOTE_ROUTE = "/response_vote"
 COLOSSEUM_ENERGY_VOTE_ROUTE = "/energy_vote"
 COLOSSEUM_HEALTH_ROUTE = "/health"
 
 
+class ModelsResponse(BaseModel):
+    available_models: list[str]
+
+
 class PromptRequest(BaseModel):
     request_id: str
     prompt: str
     model_index: Literal[0, 1]
+    model_preference: str
 
 
 class ResponseVoteRequest(BaseModel):

spitfight/colosseum/controller/controller.py

@@ -12,7 +12,7 @@ from pydantic import BaseModel, Field
 from spitfight.log import get_logger
 from spitfight.utils import BoundedExpiringDict, TokenGenerationBuffer, create_task
 from spitfight.colosseum.controller.worker import WorkerService
-from spitfight.prompt import get_system_prompt, apply_model_characteristics
+from spitfight.prompt import apply_model_characteristics
 
 if TYPE_CHECKING:
     from spitfight.colosseum.controller.router import ControllerConfig
@@ -46,6 +46,7 @@ class RequestState(BaseModel):
     request_id: str
     model_names: list[str]
     raw_prompt: str
+    model_preference: str
     responses: list[str] = ["UNSET", "UNSET"]
     model_prompts: list[str] = ["UNSET", "UNSET"]
     energy_consumptions: list[float] = [0.0, 0.0]
@@ -140,6 +141,14 @@ class Controller:
                 prev_num_req_states - len(self.request_states),
             )
 
+    def get_available_models(self) -> list[str]:
+        """Return the names of available models."""
+        return [
+            worker.model_name
+            for worker in self.worker_service.workers
+            if worker.status == "up"
+        ]
+
     def response_vote(self, request_id: str, victory_index: Literal[0, 1]) -> RequestState | None:
         """Record the user's response vote and return the new state."""
         if (state := self.request_states.get(request_id)) is not None:
@@ -165,16 +174,18 @@ class Controller:
         request_id: str,
         prompt: str,
         model_index: Literal[0, 1],
+        model_preference: str,
     ) -> AsyncGenerator[bytes, None]:
         # This method is called twice for the same request, once for each model.
         # If it's the first time this method is called, assign models to the request.
         if request_id not in self.request_states:
-            workers = self.worker_service.choose_two()
+            workers = self.worker_service.choose_based_on_preference(model_preference)
             model_names = [worker.model_name for worker in workers]
             self.request_states[request_id] = RequestState(
                 request_id=request_id,
                 raw_prompt=prompt,
                 model_names=model_names,
+                model_preference=model_preference,
             )
         request_state = self.request_states[request_id]
         model_name = request_state.model_names[model_index]

spitfight/colosseum/controller/router.py

@@ -10,10 +10,12 @@ from text_generation.errors import OverloadedError, UnknownError, ValidationErro
 
 from spitfight.log import get_logger, init_queued_root_logger, shutdown_queued_root_loggers
 from spitfight.colosseum.common import (
+    COLOSSEUM_MODELS_ROUTE,
     COLOSSEUM_PROMPT_ROUTE,
     COLOSSEUM_RESP_VOTE_ROUTE,
     COLOSSEUM_ENERGY_VOTE_ROUTE,
     COLOSSEUM_HEALTH_ROUTE,
+    ModelsResponse,
     PromptRequest,
     ResponseVoteRequest,
     ResponseVoteResponse,
@@ -67,12 +69,21 @@ async def shutdown_event():
     get_global_controller().shutdown()
     shutdown_queued_root_loggers()
 
+@app.get(COLOSSEUM_MODELS_ROUTE, response_model=ModelsResponse)
+async def models(controller: Controller = Depends(get_global_controller)):
+    return ModelsResponse(available_models=controller.get_available_models())
+
 @app.post(COLOSSEUM_PROMPT_ROUTE)
 async def prompt(
     request: PromptRequest,
     controller: Controller = Depends(get_global_controller),
 ):
-    generator = controller.prompt(request.request_id, request.prompt, request.model_index)
+    generator = controller.prompt(
+        request.request_id,
+        request.prompt,
+        request.model_index,
+        request.model_preference,
+    )
 
     # First try to get the first token in order to catch TGI errors.
     try:

spitfight/colosseum/controller/worker.py

@@ -19,7 +19,7 @@ class Worker(BaseModel):
     hostname: str
     # For TGI, this would always be 80.
     port: int
-    # User-friendly model name, e.g. "metaai/llama2-13b-chat".
+    # User-friendly model name, e.g. "Llama2-7B".
     model_name: str
     # Hugging Face model ID, e.g. "metaai/Llama-2-13b-chat-hf".
     model_id: str
@@ -146,6 +146,21 @@ class WorkerService:
         worker_a, worker_b = random.sample(live_workers, 2)
         return worker_a, worker_b
 
+    def choose_based_on_preference(self, preference: str) -> tuple[Worker, Worker]:
+        """Choose two different workers based on user preference.
+
+        Specifically, if `preference` is `"Random"`, this is equivalent to
+        choosing two models at random. Otherwise, if `preference` is a model
+        name, this is equivalent to choosing that model and another model at
+        random. In that case, the order of the two models is also randomized.
+        """
+        if preference == "Random":
+            return self.choose_two()
+        else:
+            worker_a = self.get_worker(preference)
+            worker_b = random.choice([worker for worker in self.workers if worker != worker_a])
+            return tuple(random.sample([worker_a, worker_b], 2))
+
     async def check_workers(self) -> None:
         """Check the status of all workers."""
         await asyncio.gather(*[worker.check_status() for worker in self.workers])