yzhuang/MixtureOfInputs

1from __future__ import annotations2 3import os4import openai5import gradio as gr6import server7 8# ──────────────────────────────────────────────────────────────────────────────9# OpenAI client configuration10# ──────────────────────────────────────────────────────────────────────────────11# ``openai`` still expects an API key even if the backend ignores it, so we use12# a dummy value when none is provided.  The *base_url* points to the local13# vLLM server that speaks the OpenAI REST dialect.14# -----------------------------------------------------------------------------15openai_api_key = "EMPTY"16openai_api_base = "http://0.0.0.0:8000/v1"17 18client = openai.OpenAI(19    api_key=openai_api_key,20    base_url=openai_api_base,21)22 23# ──────────────────────────────────────────────────────────────────────────────24# Chat handler25# ──────────────────────────────────────────────────────────────────────────────26 27def stream_completion(message: str,28                      history: list[tuple[str, str]],29                      max_tokens: int,30                      temperature: float,31                      top_p: float,32                      beta: float):33    """Gradio callback that yields streaming assistant replies.34 35    The function reconstructs the conversation *excluding* any system prompt36    and then calls ``openai.chat.completions.create`` with ``stream=True``.37    Each incoming delta is appended to an ``assistant`` buffer which is sent38    back to the Chatbot component for real‑time display.39    """40 41    # Build OpenAI‑style message list from prior turns42    messages: list[dict[str, str]] = []43    for user_msg, assistant_msg in history:44        if user_msg:45            messages.append({"role": "user", "content": user_msg})46        if assistant_msg:47            messages.append({"role": "assistant", "content": assistant_msg})48 49    # Current user input comes last50    messages.append({"role": "user", "content": message})51 52    os.environ["MIXINPUTS_BETA"] = str(beta)53 54    #try:55    # Kick off streaming completion56    response = client.chat.completions.create(57        model="Qwen/Qwen3-4B",58        messages=messages,59        temperature=temperature,60        top_p=top_p,61        max_tokens=max_tokens,62    )63 64    assistant = response.choices[0].message.content65    yield history + [(message, assistant)]  # live update66 67 68# ──────────────────────────────────────────────────────────────────────────────69# Gradio UI70# ──────────────────────────────────────────────────────────────────────────────71 72with gr.Blocks(title="🎨 Mixture of Inputs (MoI) Demo") as demo:73    gr.Markdown(74        "## 🎨 Mixture of Inputs (MoI) Demo with Qwen3-4B\n"75        "Streaming vLLM demo with dynamic **beta** adjustment in MoI, feel how it affects the model!\n"76        "(higher beta → less blending).\n"77        "📕Paper: https://arxiv.org/abs/2505.14827 \n"78        "💻Code: https://github.com/EvanZhuang/mixinputs \n"79    )80 81    with gr.Row():  # sliders first82        beta        = gr.Slider(0.0, 10.0, value=1.0,  step=0.1,  label="MoI β")83        temperature = gr.Slider(0.1, 1.0,  value=0.6,  step=0.1,  label="Temperature")84        top_p       = gr.Slider(0.1, 1.0,  value=0.80, step=0.05, label="Top‑p")85        max_tokens  = gr.Slider(1,   3072, value=2048,  step=1,    label="Max new tokens")86 87    chatbot   = gr.Chatbot(height=450)88    user_box  = gr.Textbox(placeholder="Type a message and press Enter…", show_label=False)89    clear_btn = gr.Button("Clear chat")90 91    user_box.submit(92        fn=stream_completion,93        inputs=[user_box, chatbot, max_tokens, temperature, top_p, beta],94        outputs=chatbot,95    )96 97    clear_btn.click(lambda: None, None, chatbot, queue=False)98 99# ──────────────────────────────────────────────────────────────────────────────100if __name__ == "__main__":101    demo.launch()102