Spaces:
Running
Running
| from random import sample | |
| from fastchat.serve import gradio_web_server | |
| from fastchat.serve.gradio_web_server import bot_response | |
| import gradio as gr | |
| # TODO(#1): Add more models. | |
| SUPPORTED_MODELS = ["gpt-4", "gpt-4-turbo", "gpt-3.5-turbo"] | |
| def user(user_message): | |
| model_pair = sample(SUPPORTED_MODELS, 2) | |
| new_state_a = gradio_web_server.State(model_pair[0]) | |
| new_state_b = gradio_web_server.State(model_pair[1]) | |
| for state in [new_state_a, new_state_b]: | |
| state.conv.append_message(state.conv.roles[0], user_message) | |
| state.conv.append_message(state.conv.roles[1], None) | |
| state.skip_next = False | |
| return [ | |
| new_state_a, new_state_b, new_state_a.model_name, new_state_b.model_name | |
| ] | |
| def bot(state_a, state_b, request: gr.Request): | |
| if state_a is None or state_b is None: | |
| raise RuntimeError(f"states cannot be None, got [{state_a}, {state_b}]") | |
| new_states = [state_a, state_b] | |
| generators = [] | |
| for state in new_states: | |
| try: | |
| # TODO(#1): Allow user to set configuration. | |
| # bot_response returns a generator yielding states and chatbots. | |
| generator = bot_response(state, | |
| temperature=0.9, | |
| top_p=0.9, | |
| max_new_tokens=100, | |
| request=request) | |
| generators.append(generator) | |
| # TODO(#1): Narrow down the exception type. | |
| except Exception as e: # pylint: disable=broad-except | |
| print(f"Error in bot_response: {e}") | |
| raise e | |
| new_responses = [None, None] | |
| # It simulates concurrent response generation from two models. | |
| while True: | |
| stop = True | |
| for i in range(2): | |
| try: | |
| generator = next(generators[i]) | |
| new_state = generator[0] | |
| new_states[i] = new_state | |
| # conv.messages is a list of [role, message]. | |
| new_responses[i] = new_state.conv.messages[-1][-1] | |
| stop = False | |
| except StopIteration: | |
| pass | |
| yield new_states + new_responses | |
| if stop: | |
| break | |
| with gr.Blocks() as app: | |
| model_names = [gr.State(None), gr.State(None)] | |
| responses = [gr.State(None), gr.State(None)] | |
| # states stores FastChat-specific conversation states. | |
| states = [gr.State(None), gr.State(None)] | |
| prompt = gr.TextArea(label="Prompt", lines=4) | |
| submit = gr.Button() | |
| with gr.Row(): | |
| responses[0] = gr.Textbox(label="Model A", interactive=False) | |
| responses[1] = gr.Textbox(label="Model B", interactive=False) | |
| with gr.Accordion("Show models", open=False): | |
| with gr.Row(): | |
| model_names[0] = gr.Textbox(label="Model A", interactive=False) | |
| model_names[1] = gr.Textbox(label="Model B", interactive=False) | |
| submit.click(user, prompt, states + model_names, | |
| queue=False).then(bot, states, states + responses) | |
| if __name__ == "__main__": | |
| # We need to enable queue to use generators. | |
| app.queue() | |
| app.launch(debug=True) | |