tiantian-paris/FRM_Study_chatbot

1#https://huggingface.co/spaces/vonliechti/SQuAD_Agent_Experiment/blob/main/app.py2import gradio as gr3from gradio import ChatMessage4from utils import stream_from_transformers_agent5from gradio.context import Context6from gradio import Request7import pickle8import os9from dotenv import load_dotenv10from agent import get_agent, DEFAULT_TASK_SOLVING_TOOLBOX11from transformers.agents import (12    DuckDuckGoSearchTool,13    ImageQuestionAnsweringTool,14    VisitWebpageTool,15)16from tools.text_to_image import TextToImageTool17from PIL import Image18from transformers import load_tool19from prompts import (20    DEFAULT_SQUAD_REACT_CODE_SYSTEM_PROMPT,21    FOCUSED_SQUAD_REACT_CODE_SYSTEM_PROMPT,22)23from pygments.formatters import HtmlFormatter24 25 26load_dotenv()27 28SESSION_PERSISTENCE_ENABLED = os.getenv("SESSION_PERSISTENCE_ENABLED", False)29 30sessions_path = "sessions.pkl"31sessions = (32    pickle.load(open(sessions_path, "rb"))33    if SESSION_PERSISTENCE_ENABLED and os.path.exists(sessions_path)34    else {}35)36 37# If currently hosted on HuggingFace Spaces, use the default model, otherwise use the local model38model_name = (39    "meta-llama/Meta-Llama-3.1-8B-Instruct"40    if os.getenv("SPACE_ID") is not None41    else "http://localhost:1234/v1"42)43 44"""45The ImageQuestionAnsweringTool from Transformers Agents 2.0 has a bug where 46it said it accepts the path to an image, but it does not. 47This class uses the adapter pattern to fix the issue, in a way that may be 48compatible with future versions of the tool even if the bug is fixed.49"""50class FixImageQuestionAnsweringTool(ImageQuestionAnsweringTool):51    def __init__(self, *args, **kwargs):52        super().__init__(*args, **kwargs)53 54    def encode(self, image: "Image | str", question: str):55        if isinstance(image, str):56            image = Image.open(image)57        return super().encode(image, question)58 59"""60The app version of the agent has access to additional tools that are not available61during benchmarking. We chose this approach to focus benchmarking on the agent's62ability to solve questions about the SQuAD dataset, without the help of general 63knowledge available on the web.  For the purposes of the project, the demo 64app has access to additional tools to provide a more interactive and engaging experience.65"""66ADDITIONAL_TOOLS = [67    DuckDuckGoSearchTool(),68    VisitWebpageTool(),69    FixImageQuestionAnsweringTool(),70    load_tool("speech_to_text"),71    load_tool("text_to_speech"),72    load_tool("translation"),73    TextToImageTool(),74]75 76# Add image tools to the default task solving toolbox, for a more visually interactive experience77TASK_SOLVING_TOOLBOX = DEFAULT_TASK_SOLVING_TOOLBOX + ADDITIONAL_TOOLS78 79# Using the focused prompt, which was the top-performing prompt during benchmarking80system_prompt = FOCUSED_SQUAD_REACT_CODE_SYSTEM_PROMPT81 82agent = get_agent(83    model_name=model_name,84    toolbox=TASK_SOLVING_TOOLBOX,85    system_prompt=system_prompt,86    use_openai=True,  # Use OpenAI instead of a local or HF model as the base LLM engine87)88 89def append_example_message(x: gr.SelectData, messages):90    if x.value["text"] is not None:91        message = x.value["text"]92    if "files" in x.value:93        if isinstance(x.value["files"], list):94            message = "Here are the files: "95            for file in x.value["files"]:96                message += f"{file}, "97        else:98            message = x.value["files"]99    messages.append(ChatMessage(role="user", content=message))100    return messages101 102 103def add_message(message, messages):104    messages.append(ChatMessage(role="user", content=message))105    return messages106 107 108def interact_with_agent(messages, request: Request):109    session_hash = request.session_hash110    prompt = messages[-1]["content"]111    agent.logs = sessions.get(session_hash + "_logs", [])112    yield messages, gr.update(113        value="<center><h1>Thinking...</h1></center>", visible=True114    )115    for msg in stream_from_transformers_agent(agent, prompt):116        if isinstance(msg, ChatMessage):117            messages.append(msg)118            yield messages, gr.update(visible=True)119        else:120            yield messages, gr.update(121                value=f"<center><h1>{msg}</h1></center>", visible=True122            )123    yield messages, gr.update(value="<center><h1>Idle</h1></center>", visible=False)124 125 126def persist(component):127 128    def resume_session(value, request: Request):129        session_hash = request.session_hash130        print(f"Resuming session for {session_hash}")131        state = sessions.get(session_hash, value)132        agent.logs = sessions.get(session_hash + "_logs", [])133        return state134 135    def update_session(value, request: Request):136        session_hash = request.session_hash137        print(f"Updating persisted session state for {session_hash}")138        sessions[session_hash] = value139        sessions[session_hash + "_logs"] = agent.logs140        if SESSION_PERSISTENCE_ENABLED:141            pickle.dump(sessions, open(sessions_path, "wb"))142 143    Context.root_block.load(resume_session, inputs=[component], outputs=component)144    component.change(update_session, inputs=[component], outputs=None)145 146    return component147 148 149from gradio.components import (150    Component as GradioComponent,151)152from gradio.components.chatbot import (153    Chatbot,154    FileDataDict,155    FileData,156    ComponentMessage,157    FileMessage,158)159 160 161class CleanChatBot(Chatbot):162    def __init__(self, **kwargs):163        super().__init__(**kwargs)164 165    def _postprocess_content(166        self,167        chat_message: (168            str | tuple | list | FileDataDict | FileData | GradioComponent | None169        ),170    ) -> str | FileMessage | ComponentMessage | None:171        response = super()._postprocess_content(chat_message)172        print(f"Post processing content: {response}")173        if isinstance(response, ComponentMessage):174            print(f"Setting open to False for {response}")175            response.props["open"] = False176        return response177 178 179with gr.Blocks(180    fill_height=True,181    css=".gradio-container .message .content {text-align: left;}"182    + HtmlFormatter().get_style_defs(".highlight"),183) as demo:184    state = gr.State()185    inner_monologue_component = gr.Markdown(186        """<h2>Inner Monologue</h2>""", visible=False187    )188    chatbot = persist(189        gr.Chatbot(190            value=[],191            label="SQuAD Agent",192            type="messages",193            avatar_images=(194                None,195                "SQuAD.png",196            ),197            scale=1,198            autoscroll=True,199            show_copy_all_button=True,200            show_copy_button=True,201            placeholder="""<h1>SQuAD Agent</h1>202            <h2>I am your friendly guide to the Stanford Question and Answer Dataset (SQuAD).</h2>203            <h2>You can ask me questions about the dataset. You can also ask me to create images 204            to help illustrate the topics under discussion, or expand the discussion beyond the dataset.</h2>205        """,206            examples=[207                {208                    "text": "What is on top of the Notre Dame building?",209                },210                {211                    "text": "What is the Olympic Torch made of?",212                },213                {214                    "text": "Draw a picture of whatever is on top of the Notre Dame building.",215                },216            ],217        )218    )219    text_input = gr.Textbox(lines=1, label="Chat Message", scale=0)220    chat_msg = text_input.submit(add_message, [text_input, chatbot], [chatbot])221    bot_msg = chat_msg.then(222        interact_with_agent, [chatbot], [chatbot, inner_monologue_component]223    )224    text_input.submit(lambda: "", None, text_input)225    chatbot.example_select(append_example_message, [chatbot], [chatbot]).then(226        interact_with_agent, [chatbot], [chatbot, inner_monologue_component]227    )228 229if __name__ == "__main__":230    demo.launch()