web_fu_demo.py

import json
import torch
import streamlit as st
import openai

openai.api_key = "EMPTY"
openai.base_url = "http://localhost:7000/v1/"
model = "lora_model"

st.set_page_config(page_title="FOllOW-UP")
st.title("MEDICAL FOllOW-UP")


def clear_chat_history():
    del st.session_state.messages


def init_chat_history():
    with st.chat_message("assistant", avatar="🤖"):
        st.markdown("您好，我是东大FOLLOW-UP小助手，很高兴为您服务🥰")

    if "messages" in st.session_state:
        for message in st.session_state.messages:
            avatar = "🧑‍💻" if message["role"] == "user" else "🤖"
            with st.chat_message(message["role"], avatar=avatar):
                st.markdown(message["content"])
    else:
        st.session_state.messages = []

    return st.session_state.messages


@st.cache_resource
def get_re_question(messages):
    completion = openai.chat.completions.create(model=model, messages=messages)
    return completion.choices[0].message.content


def main():
    messages = init_chat_history()

    if prompt := st.chat_input("Shift + Enter 换行, Enter 发送"):
        with st.chat_message("user", avatar='🧑‍💻'):
            st.markdown(prompt)
        messages.append({"role": "user", "content": prompt})
        print(f"[user] {prompt}", flush=True)
        with st.chat_message("assistant", avatar='🤖'):
            placeholder = st.empty()
            response = get_re_question(messages)
            placeholder.markdown(response)
            if torch.backends.mps.is_available():
                torch.mps.empty_cache()
        messages.append({"role": "assistant", "content": response})
        print(json.dumps(messages, ensure_ascii=False), flush=True)

        st.button("清空对话", on_click=clear_chat_history)


if __name__ == "__main__":
    main()