[Feat]: Add streamlit UI that supports multimodal inputs (#35)

* [Feat]: Add streamlit UI that supports multimodal inputs update docs fix keys * clear history botton, and fix meta template * update demo image * fix comment --------- Co-authored-by: liukuikun <641417025@qq.com>
2023-09-20 17:08:42 +08:00
parent f05069226e
commit f20ddf3f54
7 changed files with 270 additions and 8 deletions
--- a/README.md
+++ b/README.md
@@ -1,8 +1,19 @@
-# Lagent: A lightweight framework for building LLM-based agents
+<div align="center">
+  <img src="docs/imgs/lagent_logo.png" width="450"/>
+
+[![docs](https://img.shields.io/badge/docs-latest-blue)](https://lagent.readthedocs.io/en/latest/)
+[![PyPI](https://img.shields.io/pypi/v/lagent)](https://pypi.org/project/lagent)
+[![license](https://img.shields.io/github/license/InternLM/lagent.svg)](https://github.com/InternLM/lagent/tree/main/LICENSE)
+[![issue resolution](https://img.shields.io/github/issues-closed-raw/InternLM/lagent)](https://github.com/InternLM/lagent/issues)
+[![open issues](https://img.shields.io/github/issues-raw/InternLM/lagent)](https://github.com/InternLM/lagent/issues)

 English | [简体中文](README_zh-CN.md)

-## Introduction
+</div>
+
+<p align="center">
+    👋 join us on <a href="https://twitter.com/intern_lm" target="_blank">Twitter</a>, <a href="https://discord.gg/xa29JuW87d" target="_blank">Discord</a> and <a href="https://r.vansin.top/?r=internwx" target="_blank">WeChat</a>
+</p>## Introduction

 Lagent is a lightweight open-source framework that allows users to efficiently build large language model(LLM)-based agents. It also provides some typical tools to augment LLM. The overview of our framework is shown below:

@@ -36,6 +47,17 @@ cd lagent
 pip install -e .
 ```

+### Run ReAct Web Demo
+
+```bash
+# You need to install streamlit first
+# pip install streamlit
+streamlit run examples/react_web_demo.py
+```
+
+Then you can chat through the UI shown as below
+![image](https://github.com/InternLM/lagent/assets/24622904/3aebb8b4-07d1-42a2-9da3-46080c556f68)
+
 ### Run a ReWOO agent with GPT-3.5

 Below is an example for running ReWOO with GPT-3.5
--- a/README_zh-CN.md
+++ b/README_zh-CN.md
@@ -1,4 +1,19 @@
-# Lagent: A lightweight framework for building LLM-based agents
+<div align="center">
+  <img src="docs/imgs/lagent_logo.png" width="450"/>
+
+[![docs](https://img.shields.io/badge/docs-latest-blue)](https://lagent.readthedocs.io/en/latest/)
+[![PyPI](https://img.shields.io/pypi/v/lagent)](https://pypi.org/project/lagent)
+[![license](https://img.shields.io/github/license/InternLM/lagent.svg)](https://github.com/InternLM/lagent/tree/main/LICENSE)
+[![issue resolution](https://img.shields.io/github/issues-closed-raw/InternLM/lagent)](https://github.com/InternLM/lagent/issues)
+[![open issues](https://img.shields.io/github/issues-raw/InternLM/lagent)](https://github.com/InternLM/lagent/issues)
+
+[English](README.md) | 简体中文
+
+</div>
+
+<p align="center">
+    👋 join us on <a href="https://twitter.com/intern_lm" target="_blank">Twitter</a>, <a href="https://discord.gg/xa29JuW87d" target="_blank">Discord</a> and <a href="https://r.vansin.top/?r=internwx" target="_blank">WeChat</a>
+</p>

 [English](README.md) | 简体中文

@@ -34,6 +49,18 @@ cd lagent
 pip install -e .
 ```

+### 运行一个 ReAct 智能体的网页样例
+
+```bash
+# 可能先需要安装 streamlit 包
+# pip install streamlit
+streamlit run examples/react_web_demo.py
+```
+
+然后你就可以在网页端和智能体进行对话了，效果如下图所示
+
+![image](https://github.com/InternLM/lagent/assets/24622904/3aebb8b4-07d1-42a2-9da3-46080c556f68)
+
 ### 用 GPT-3.5 构建一个 ReWOO 智能体

 ```python
--- a/docs/imgs/lagent_icon.png
+++ b/docs/imgs/lagent_icon.png
--- a/docs/imgs/lagent_logo.png
+++ b/docs/imgs/lagent_logo.png
--- a/examples/react_web_demo.py
+++ b/examples/react_web_demo.py
@@ -0,0 +1,215 @@
+import copy
+import os
+
+import streamlit as st
+from streamlit.logger import get_logger
+
+from lagent.actions import ActionExecutor, GoogleSearch, PythonInterpreter
+from lagent.agents.react import ReAct
+from lagent.llms import GPTAPI
+from lagent.llms.huggingface import HFTransformerCasualLM
+
+
+class SessionState:
+
+    def init_state(self):
+        """Initialize session state variables."""
+        st.session_state['assistant'] = []
+        st.session_state['user'] = []
+
+        action_list = [PythonInterpreter(), GoogleSearch()]
+        st.session_state['plugin_map'] = {
+            action.name: action
+            for action in action_list
+        }
+        st.session_state['model_map'] = {}
+        st.session_state['model_selected'] = None
+        st.session_state['plugin_actions'] = set()
+
+    def clear_state(self):
+        """Clear the existing session state."""
+        st.session_state['assistant'] = []
+        st.session_state['user'] = []
+        st.session_state['model_selected'] = None
+        if 'chatbot' in st.session_state:
+            st.session_state['chatbot']._session_history = []
+
+
+class StreamlitUI:
+
+    def __init__(self, session_state: SessionState):
+        self.init_streamlit()
+        self.session_state = session_state
+
+    def init_streamlit(self):
+        """Initialize Streamlit's UI settings."""
+        st.set_page_config(
+            layout='wide',
+            page_title='lagent-web',
+            page_icon='./docs/imgs/lagent_icon.png')
+        st.header(':robot_face: :blue[Lagent] Web Demo ', divider='rainbow')
+        st.sidebar.title('模型控制')
+
+    def setup_sidebar(self):
+        """Setup the sidebar for model and plugin selection."""
+        model_name = st.sidebar.selectbox(
+            '模型选择：', options=['gpt-3.5-turbo', 'internlm'])
+        if model_name != st.session_state['model_selected']:
+            model = self.init_model(model_name)
+            self.session_state.clear_state()
+            st.session_state['model_selected'] = model_name
+            if 'chatbot' in st.session_state:
+                del st.session_state['chatbot']
+        else:
+            model = st.session_state['model_map'][model_name]
+
+        plugin_name = st.sidebar.multiselect(
+            '插件选择',
+            options=list(st.session_state['plugin_map'].keys()),
+            default=[list(st.session_state['plugin_map'].keys())[0]],
+        )
+
+        plugin_action = [
+            st.session_state['plugin_map'][name] for name in plugin_name
+        ]
+        if 'chatbot' in st.session_state:
+            st.session_state['chatbot']._action_executor = ActionExecutor(
+                actions=plugin_action)
+        if st.sidebar.button('清空对话', key='clear'):
+            self.session_state.clear_state()
+        uploaded_file = st.sidebar.file_uploader(
+            '上传文件', type=['png', 'jpg', 'jpeg', 'mp4', 'mp3', 'wav'])
+        return model_name, model, plugin_action, uploaded_file
+
+    def init_model(self, option):
+        """Initialize the model based on the selected option."""
+        if option not in st.session_state['model_map']:
+            if option.startswith('gpt'):
+                st.session_state['model_map'][option] = GPTAPI(
+                    model_type=option)
+            else:
+                st.session_state['model_map'][option] = HFTransformerCasualLM(
+                    'internlm/internlm-chat-7b-v1_1')
+        return st.session_state['model_map'][option]
+
+    def initialize_chatbot(self, model, plugin_action):
+        """Initialize the chatbot with the given model and plugin actions."""
+        return ReAct(
+            llm=model, action_executor=ActionExecutor(actions=plugin_action))
+
+    def render_user(self, prompt: str):
+        with st.chat_message('user'):
+            st.markdown(prompt)
+
+    def render_assistant(self, agent_return):
+        with st.chat_message('assistant'):
+            for action in agent_return.actions:
+                if (action):
+                    self.render_action(action)
+            st.markdown(agent_return.response)
+
+    def render_action(self, action):
+        with st.expander(action.type, expanded=True):
+            st.markdown(
+                "<p style='text-align: left;display:flex;'> <span style='font-size:14px;font-weight:600;width:70px;text-align-last: justify;'>插    件</span><span style='width:14px;text-align:left;display:block;'>:</span><span style='flex:1;'>"  # noqa E501
+                + action.type + '</span></p>',
+                unsafe_allow_html=True)
+            st.markdown(
+                "<p style='text-align: left;display:flex;'> <span style='font-size:14px;font-weight:600;width:70px;text-align-last: justify;'>思考步骤</span><span style='width:14px;text-align:left;display:block;'>:</span><span style='flex:1;'>"  # noqa E501
+                + action.thought + '</span></p>',
+                unsafe_allow_html=True)
+            if (isinstance(action.args, dict) and 'text' in action.args):
+                st.markdown(
+                    "<p style='text-align: left;display:flex;'><span style='font-size:14px;font-weight:600;width:70px;text-align-last: justify;'> 执行内容</span><span style='width:14px;text-align:left;display:block;'>:</span></p>",  # noqa E501
+                    unsafe_allow_html=True)
+                st.markdown(action.args['text'])
+            self.render_action_results(action)
+
+    def render_action_results(self, action):
+        """Render the results of action, including text, images, videos, and
+        audios."""
+        if (isinstance(action.result, dict)):
+            st.markdown(
+                "<p style='text-align: left;display:flex;'><span style='font-size:14px;font-weight:600;width:70px;text-align-last: justify;'> 执行结果</span><span style='width:14px;text-align:left;display:block;'>:</span></p>",  # noqa E501
+                unsafe_allow_html=True)
+            if 'text' in action.result:
+                st.markdown(
+                    "<p style='text-align: left;'>" + action.result['text'] +
+                    '</p>',
+                    unsafe_allow_html=True)
+            if 'image' in action.result:
+                image_path = action.result['image']
+                image_data = open(image_path, 'rb').read()
+                st.image(image_data, caption='Generated Image')
+            if 'video' in action.result:
+                video_data = action.result['video']
+                video_data = open(video_data, 'rb').read()
+                st.video(video_data)
+            if 'audio' in action.result:
+                audio_data = action.result['audio']
+                audio_data = open(audio_data, 'rb').read()
+                st.audio(audio_data)
+
+
+def main():
+    logger = get_logger(__name__)
+    # Initialize Streamlit UI and setup sidebar
+    if 'ui' not in st.session_state:
+        session_state = SessionState()
+        session_state.init_state()
+        st.session_state['ui'] = StreamlitUI(session_state)
+
+    else:
+        st.set_page_config(
+            layout='wide',
+            page_title='lagent-web',
+            page_icon='./docs/imgs/lagent_icon.png')
+        st.header(':robot_face: :blue[Lagent] Web Demo ', divider='rainbow')
+    model_name, model, plugin_action, uploaded_file = st.session_state[
+        'ui'].setup_sidebar()
+
+    # Initialize chatbot if it is not already initialized
+    # or if the model has changed
+    if 'chatbot' not in st.session_state or model != st.session_state[
+            'chatbot']._llm:
+        st.session_state['chatbot'] = st.session_state[
+            'ui'].initialize_chatbot(model, plugin_action)
+
+    for prompt, agent_return in zip(st.session_state['user'],
+                                    st.session_state['assistant']):
+        st.session_state['ui'].render_user(prompt)
+        st.session_state['ui'].render_assistant(agent_return)
+    # User input form at the bottom (this part will be at the bottom)
+    # with st.form(key='my_form', clear_on_submit=True):
+
+    if user_input := st.chat_input(''):
+        st.session_state['ui'].render_user(user_input)
+        st.session_state['user'].append(user_input)
+        # Add file uploader to sidebar
+        if uploaded_file:
+            file_bytes = uploaded_file.read()
+            file_type = uploaded_file.type
+            if 'image' in file_type:
+                st.image(file_bytes, caption='Uploaded Image')
+            elif 'video' in file_type:
+                st.video(file_bytes, caption='Uploaded Video')
+            elif 'audio' in file_type:
+                st.audio(file_bytes, caption='Uploaded Audio')
+            # Save the file to a temporary location and get the path
+            file_path = os.path.join(root_dir, uploaded_file.name)
+            with open(file_path, 'wb') as tmpfile:
+                tmpfile.write(file_bytes)
+            st.write(f'File saved at: {file_path}')
+            user_input = '我上传了一个图像，路径为: {file_path}. {user_input}'.format(
+                file_path=file_path, user_input=user_input)
+        agent_return = st.session_state['chatbot'].chat(user_input)
+        st.session_state['assistant'].append(copy.deepcopy(agent_return))
+        logger.info(agent_return.inner_steps)
+        st.session_state['ui'].render_assistant(agent_return)
+
+
+if __name__ == '__main__':
+    root_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    root_dir = os.path.join(root_dir, 'tmp_dir')
+    os.makedirs(root_dir, exist_ok=True)
+    main()
--- a/lagent/llms/huggingface.py
+++ b/lagent/llms/huggingface.py
@@ -47,11 +47,8 @@ class HFTransformer(BaseModel):
            tokenizer_only: bool = False,
            model_kwargs: dict = dict(device_map='auto'),
            meta_template: Optional[Dict] = [
-                dict(
-                    role='system',
-                    begin='<|System|>:',
-                    end='<TOKENS_UNUSED_2>\n'),
-                dict(role='user', begin='<|User|>:', end='<eoh>\n'),
+                dict(role='system', begin='<|System|>:', end='\n'),
+                dict(role='user', begin='<|User|>:', end='\n'),
                dict(
                    role='assistant',
                    begin='<|Bot|>:',
--- a/requirements/optional.txt
+++ b/requirements/optional.txt
@@ -1,2 +1,3 @@
 torch
 transformers
+streamlit