pull/599/head
gaoyang07 2024-01-17 13:01:56 +08:00
parent feeb3efbc5
commit 9e56989066
8 changed files with 19 additions and 14 deletions

View File

@ -11,11 +11,10 @@ from dataclasses import asdict
import streamlit as st import streamlit as st
import torch import torch
from tools.transformers.interface import GenerationConfig, generate_interactive
from transformers import AutoModelForCausalLM, AutoTokenizer from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.utils import logging from transformers.utils import logging
from tools.transformers.interface import GenerationConfig, generate_interactive
logger = logging.get_logger(__name__) logger = logging.get_logger(__name__)
@ -109,9 +108,15 @@ def main():
): ):
# Display robot response in chat message container # Display robot response in chat message container
message_placeholder.markdown(cur_response + "") message_placeholder.markdown(cur_response + "")
message_placeholder.markdown(cur_response) message_placeholder.markdown(cur_response) # pylint: disable=undefined-loop-variable
# Add robot response to chat history # Add robot response to chat history
st.session_state.messages.append({"role": "robot", "content": cur_response, "avatar": robot_avator}) st.session_state.messages.append(
{
"role": "robot",
"content": cur_response, # pylint: disable=undefined-loop-variable
"avatar": robot_avator,
}
)
torch.cuda.empty_cache() torch.cuda.empty_cache()