StewartLab/lm-survey-interface

1"""2General-Purpose LM Interview Interface3 4Author: Dr Musashi Hinck5 6 7Version Log:8 9- 2024.01.29: prototype without separate launching interface for demoing in SPIA class.10    - Remove URL decoding11    - Read sysprompt and initial_message from file12    - Begins with user entering name/alias13    - Azure OpenAI?14- 2024.01.31: wandb does not work for use case, what to do instead?15    - Write to local file and then upload at end? (does filestream cause blocking?)16- 2024.03.03: Creating new instance for demoing to IRB17 18"""19from __future__ import annotations20 21import os22import logging23import json24import wandb25import gradio as gr26from typing import Generator, Any27 28from pathlib import Path29 30logger = logging.getLogger(__name__)31 32from utils import (33    PromptTemplate,34    convert_gradio_to_openai,35    initialize_client,36    seed_azure_key37)38 39 40# %% Initialization41CONFIG_DIR: Path = Path("./CogDebIRB")42if os.environ.get("AZURE_ENDPOINT") is None: # Set Azure credentials from local files43    seed_azure_key()44client = initialize_client()45 46# %% (functions)47def load_config(48    path: Path,49) -> tuple[str, str, dict[str, str | float], dict[str, str | list[str]]]:50    "Read configs, return inital_message, system_message, model_args, wandb_args"51    initial_message: str = (path / "initial_message.txt").read_text().strip()52    system_message: str = (path / "system_message.txt").read_text().strip()53    cfg: dict[str, str] = json.loads((path / "config.json").read_bytes())54    model_args: dict[str, str | float] = cfg.get(55        "model_args", {"model": "gpt4", "temperature": 0.0}56    )57    wandb_args: dict = cfg.get("wandb_args")58    return initial_message, system_message, model_args, wandb_args59 60 61def initialize_interview(62    initial_message: str,63) -> tuple[gr.Chatbot,64           gr.Textbox,65           gr.Button,66           gr.Button,67           gr.Button]:68    "Read system prompt and start interview. Change visibilities of elements."69    chat_history = [70        [None, initial_message]71    ]  # First item is for user, in this case bot starts interaction.72    return (73        gr.Chatbot(visible=True, value=chat_history), # chatDisplay74        gr.Textbox(75            placeholder="Type response here. Hit 'enter' to submit.",76            visible=True,77            interactive=True,78        ),  # chatInput79        gr.Button(visible=True, interactive=True),  # chatSubmit80        gr.Button(visible=False),  # startInterview81        gr.Button(visible=True),  # resetButton82    )83 84 85def initialize_tracker(86    model_args: dict[str, str | float],87    system_message: PromptTemplate,88    userid: str,89    wandb_args: dict[str, str | list[str]],90) -> gr.Textbox:91    "Initializes wandb run for interview. Resets userBox afterwards."92    run_config = model_args | {93        "system_message": str(system_message),94        "userid": userid,95    }96    logger.info(f"Initializing WandB run for {userid}")97    wandb.init(98        project=wandb_args["project"],99        name=userid,100        config=run_config,101        tags=wandb_args["tags"],102    )103    return gr.Textbox(value=None, visible=False)104 105 106def save_interview(107    chat_history: list[list[str | None]],108) -> None:109    # Save chat_history as json110    with open(CONFIG_DIR/"transcript.json", 'w') as fh:111        json.dump(chat_history, fh, indent=2)112    chat_data = []113    for pair in chat_history:114        for i, role in enumerate(["user", "bot"]):115            if pair[i] is not None:116                chat_data += [[role, pair[i]]]117    chat_table = wandb.Table(data=chat_data, columns=["role", "message"])118    logger.info("Uploading interview transcript to WandB...")119    wandb.log({"chat_history": chat_table})120    logger.info("Uploading complete.")121 122 123 124def user_message(125    message: str, chat_history: list[list[str | None]]126) -> tuple[str, list[list[str | None]]]:127    "Display user message immediately"128    return "", chat_history + [[message, None]]129 130 131def bot_message(132    chat_history: list[list[str | None]],133    system_message: str,134    model_args: dict[str, str | float],135) -> Generator[Any, Any, Any]:136    # Prep messages137    user_msg = chat_history[-1][0]138    messages = convert_gradio_to_openai(chat_history[:-1])139    messages = (140        [{"role": "system", "content": system_message}]141        + messages142        + [{"role": "user", "content": user_msg}]143    )144    # API call145    response = client.chat.completions.create(146        messages=messages, stream=True, **model_args147    )148    # Streaming149    chat_history[-1][1] = ""150    for chunk in response:151        delta = chunk.choices[0].delta.content152        if delta:153            chat_history[-1][1] += delta154            yield chat_history155 156 157def reset_interview() -> (158    tuple[159        list[list[str | None]], gr.Chatbot, gr.Textbox, gr.Button, gr.Button, gr.Button160    ]161):162    wandb.finish()163    gr.Info("Interview reset.")164    return (165        gr.Chatbot(visible=False, value=[]),  # chatDisplay166        gr.Textbox(visible=False),  # chatInput167        gr.Button(visible=False),  # chatSubmit168        gr.Textbox(value=None, visible=True),  # userBox169        gr.Button(visible=True),  # startInterview170        gr.Button(visible=False),  # resetButton171    )172 173 174# LAYOUT175with gr.Blocks(theme="sudeepshouche/minimalist") as demo:176    gr.Markdown("# Chat Interview Interface")177    userDisplay = gr.Markdown("", visible=False)178 179    # Config values180    configDir = gr.State(value=CONFIG_DIR)181    initialMessage = gr.Textbox(visible=False)182    systemMessage = gr.Textbox(visible=False)183    modelArgs = gr.State(value={"model": "", "temperature": ""})184    wandbArgs = gr.State(value={"project": "", "tags": []})185 186    ## Start interview by entering name or alias187    userBox = gr.Textbox(188        value=None, placeholder="Enter name or alias and hit 'enter' to begin.", show_label=False189    )190    startInterview = gr.Button("Start Interview", variant="primary", visible=True)191 192    ## RESPONDENT193    chatDisplay = gr.Chatbot(show_label=False, visible=False)194    with gr.Row():195        chatInput = gr.Textbox(196            placeholder="Click 'Start Interview' to begin.",197            visible=False,198            interactive=False,199            show_label=False,200            scale=10,201        )202        chatSubmit = gr.Button(203            "",204            variant="primary",205            interactive=False,206            icon="./arrow_icon.svg",207            visible=False,208        )209    resetButton = gr.Button("Save and Exit", visible=False, variant="stop")210    disclaimer = gr.HTML(211        """212        <div213        style='font-size: 1em;214               font-style: italic;   215               position: fixed;216               left: 50%;217               bottom: 20px;218               transform: translate(-50%, -50%);219               margin: 0 auto;220               '221        >{}</div>222        """.format(223            "Statements by the chatbot may contain factual inaccuracies."224        )225    )226 227    ## INTERACTIONS228    # Start Interview button229    userBox.change(lambda x: x, inputs=[userBox], outputs=[userDisplay], show_progress=False)230    userBox.submit(231        load_config,232        inputs=configDir,233        outputs=[initialMessage, systemMessage, modelArgs, wandbArgs],234    ).then(235        initialize_interview,236        inputs=[initialMessage],237        outputs=[238            chatDisplay,239            chatInput,240            chatSubmit,241            startInterview,242            resetButton,243        ],244    ).then(245        initialize_tracker,246        inputs=[modelArgs, systemMessage, userBox, wandbArgs],247        outputs=[userBox]248    )249 250    startInterview.click(251        load_config,252        inputs=configDir,253        outputs=[initialMessage, systemMessage, modelArgs, wandbArgs],254    ).then(255        initialize_interview,256        inputs=[initialMessage],257        outputs=[258            chatDisplay,259            chatInput,260            chatSubmit,261            startInterview,262            resetButton,263        ],264    ).then(265        initialize_tracker,266        inputs=[modelArgs, systemMessage, userBox, wandbArgs],267        outputs=[userBox]268    )269 270    # Chat interaction271    # "Enter"272    chatInput.submit(273        user_message,274        inputs=[chatInput, chatDisplay],275        outputs=[chatInput, chatDisplay],276        queue=False,277    ).then(278        bot_message,279        inputs=[chatDisplay, systemMessage, modelArgs],280        outputs=[chatDisplay],281    ).then(282        save_interview, inputs=[chatDisplay]283    )284    # Button285    chatSubmit.click(286        user_message,287        inputs=[chatInput, chatDisplay],288        outputs=[chatInput, chatDisplay],289        queue=False,290    ).then(291        bot_message,292        inputs=[chatDisplay, systemMessage, modelArgs],293        outputs=[chatDisplay],294    ).then(295        save_interview, inputs=[chatDisplay]296    )297 298    # Reset button299    resetButton.click(save_interview, [chatDisplay]).then(300        reset_interview,301        outputs=[302            chatDisplay,303            chatInput,304            chatSubmit,305            userBox,306            startInterview,307            resetButton,308        ],309        show_progress=False,310    )311 312 313if __name__ == "__main__":314    demo.launch()315