def _launch_ui()

in model-gallery/deploy/mllm/webui_client.py [0:0]

• 164 lines of code
• 18 McCabe index (conditional complexity)

def _launch_ui(model, client, args):
    def call_local_model(messages, max_tokens, temperature, timeout, use_stream):
        try:
            messages = transform_messages(messages)
            print("Messages:", messages)
            print(
                f"Arguments: max_tokens={max_tokens}, temperature={temperature}, timeout={timeout}, stream={use_stream}"
            )
            gen = client.chat.completions.create(
                messages=messages,
                model=model,
                max_tokens=max_tokens,
                temperature=temperature,
                timeout=timeout,
                stream=use_stream,
            )
            if use_stream:
                generated_text = ""
                for chunk in gen:
                    generated_text += chunk.choices[0].delta.content
                    yield generated_text
            else:
                generated_text = gen.choices[0].message.content
                yield generated_text
        except Exception as e:
            print(e)
            raise gr.Error(e.message)

    def create_predict_fn():
        def predict(
            _chatbot, task_history, max_tokens, temperature, timeout, use_stream
        ):
            chat_query = _chatbot[-1][0]
            query = task_history[-1][0]
            if len(chat_query) == 0:
                _chatbot.pop()
                task_history.pop()
                return _chatbot
            print("User: " + _parse_text(query))
            history_cp = copy.deepcopy(task_history)
            full_response = ""
            messages = []
            content = []
            for cls, q, a in history_cp:
                if cls == "image":
                    content.append({"image_url": f"data:image/jpeg;base64,{q}"})
                elif cls == "video":
                    content.append({"video_url": f"data:video/mp4;base64,{q}"})
                else:
                    content.append({"text": q})
                    messages.append({"role": "user", "content": content})
                    messages.append({"role": "assistant", "content": [{"text": a}]})
                    content = []
            messages.pop()

            for response in call_local_model(
                messages, max_tokens, temperature, timeout, use_stream
            ):
                _chatbot[-1] = (
                    _parse_text(chat_query),
                    _remove_image_special(_parse_text(response)),
                )

                yield _chatbot
                full_response = _parse_text(response)

            task_history[-1] = ("text", query, full_response)
            print("VL-Chat: " + _parse_text(full_response))
            yield _chatbot

        return predict

    def create_regenerate_fn():
        def regenerate(
            _chatbot, task_history, max_tokens, temperature, timeout, use_stream
        ):
            print("Regenerating...")
            if not task_history:
                return _chatbot
            item = task_history[-1]
            if item[1] is None:
                return _chatbot
            task_history[-1] = (item[0], None)
            chatbot_item = _chatbot.pop(-1)
            if chatbot_item[0] is None:
                _chatbot[-1] = (_chatbot[-1][0], None)
            else:
                _chatbot.append((chatbot_item[0], None))
            _chatbot_gen = predict(
                _chatbot, task_history, max_tokens, temperature, timeout, use_stream
            )
            for _chatbot in _chatbot_gen:
                yield _chatbot

        return regenerate

    predict = create_predict_fn()
    regenerate = create_regenerate_fn()

    def add_text(history, task_history, text):
        task_text = text
        history = history if history is not None else []
        task_history = task_history if task_history is not None else []
        history = history + [(_parse_text(text), None)]
        task_history = task_history + [("text", task_text, None)]
        return history, task_history

    def add_file(history, task_history, file):
        history = history if history is not None else []
        task_history = task_history if task_history is not None else []
        base64_encoded_content = encode_base64_content_from_file(file.name)
        history = history + [((file.name,), None)]
        if is_video_file(file.name):
            task_history = task_history + [("video", base64_encoded_content, None)]
        else:
            task_history = task_history + [("image", base64_encoded_content, None)]
        return history, task_history

    def reset_user_input():
        return gr.update(value="")

    def reset_state(task_history):
        print("Clear history.")
        task_history.clear()
        return []

    with gr.Blocks() as demo:

        gr.Markdown(f"""<center><font size=8>MLLM-WebUI</center>""")
        chatbot = gr.Chatbot(elem_classes="control-height", height=480)
        query = gr.Textbox(lines=2, label="Input")
        task_history = gr.State([])

        with gr.Row():
            max_tokens = gr.Slider(
                minimum=10, maximum=10240, step=10, label="max_tokens", value=1024
            )
            temperature = gr.Slider(
                minimum=0.0, maximum=2.0, step=0.01, label="temperature", value=0.0
            )
            timeout = gr.Slider(
                minimum=0, maximum=600, step=10, label="timeout(seconds)", value=120
            )
            use_stream_chat = gr.Checkbox(label="use_stream_chat", value=True)

        with gr.Row():
            addfile_btn = gr.UploadButton(
                "📁 Upload (上传文件)", file_types=["image", "video"]
            )
            submit_btn = gr.Button("🚀 Submit (发送)")
            regen_btn = gr.Button("🤔️ Regenerate (重试)")
            empty_bin = gr.Button("🧹 Clear History (清除历史)")

        submit_btn.click(
            add_text, [chatbot, task_history, query], [chatbot, task_history]
        ).then(
            predict,
            [chatbot, task_history, max_tokens, temperature, timeout, use_stream_chat],
            [chatbot],
            show_progress=True,
        )
        submit_btn.click(reset_user_input, [], [query])
        empty_bin.click(reset_state, [task_history], [chatbot], show_progress=True)
        regen_btn.click(
            regenerate,
            [chatbot, task_history, max_tokens, temperature, timeout, use_stream_chat],
            [chatbot],
            show_progress=True,
        )
        addfile_btn.upload(
            add_file,
            [chatbot, task_history, addfile_btn],
            [chatbot, task_history],
            show_progress=True,
        )

    demo.queue().launch(
        share=args.share,
        inbrowser=args.inbrowser,
        server_port=args.server_port,
        server_name=args.server_name,
    )