Spaces:
Running
Running
| from random import sample | |
| import gradio as gr | |
| from fastchat.serve import gradio_web_server | |
| from fastchat.serve.gradio_web_server import bot_response | |
| # TODO(#1): Add more models. | |
| SUPPORTED_MODELS = ["gpt-4", "gpt-4-turbo", "gpt-3.5-turbo"] | |
| def user(state_a, state_b, user_message): | |
| if state_a is None and state_b is None: | |
| model_pair = sample(SUPPORTED_MODELS, 2) | |
| state_a = gradio_web_server.State(model_pair[0]) | |
| state_b = gradio_web_server.State(model_pair[1]) | |
| for state in [state_a, 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 | |
| empty_prompt = "" | |
| return [ | |
| state_a, state_b, | |
| state_a.to_gradio_chatbot(), | |
| state_b.to_gradio_chatbot(), state_a.model_name, state_b.model_name, | |
| empty_prompt | |
| ] | |
| 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}]") | |
| generators = [] | |
| for state in [state_a, state_b]: | |
| 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_chatbots = [None, None] | |
| while True: | |
| stop = True | |
| for i in range(2): | |
| try: | |
| generator = next(generators[i]) | |
| states[i], new_chatbots[i] = generator[0], generator[1] | |
| stop = False | |
| except StopIteration: | |
| pass | |
| yield [state_a, state_b] + new_chatbots | |
| if stop: | |
| break | |
| with gr.Blocks() as app: | |
| with gr.Row(): | |
| response_type = gr.Radio( | |
| ["Summarization", "Translation"], | |
| value="Summarization", | |
| label="Response type", | |
| info="Choose the type of response you want from the model.") | |
| language = gr.Dropdown(["Korean", "English"], | |
| value="Korean", | |
| label="Language", | |
| info="Choose the target language.") | |
| chatbots = [None, None] | |
| with gr.Row(): | |
| chatbots[0] = gr.Chatbot(label="Model A") | |
| chatbots[1] = gr.Chatbot(label="Model B") | |
| model_names = [None, None] | |
| 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) | |
| prompt = gr.Textbox(label="Prompt") | |
| states = [gr.State(None), gr.State(None)] | |
| prompt.submit(user, | |
| states + [prompt], | |
| states + chatbots + model_names + [prompt], | |
| queue=False).then(bot, states, states + chatbots) | |
| if __name__ == "__main__": | |
| # We need to enable queue to use generators. | |
| app.queue() | |
| app.launch() | |