update inference

This commit is contained in:
gongjy 2024-10-13 13:09:52 +08:00
parent 135421690e
commit ad6fc92399

View File

@ -4,15 +4,17 @@ import torch
from transformers import AutoModelForCausalLM, AutoTokenizer from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation.utils import GenerationConfig from transformers.generation.utils import GenerationConfig
st.set_page_config(page_title="MiniMind-V1 108M(无历史上文)") st.set_page_config(page_title="minimind-v1(108M)")
st.title("MiniMind-V1 108M(无历史上文)") st.title("minimind-v1(108M)")
model_id = "minimind-v1" model_id = "minimind-v1"
# ----------------------------------------------------------------------------- # -----------------------------------------------------------------------------
temperature = 0.7 temperature = 0.5
top_k = 8 top_k = 16
max_seq_len = 1 * 1024 max_seq_len = 1 * 1024
# ----------------------------------------------------------------------------- # -----------------------------------------------------------------------------
@st.cache_resource @st.cache_resource
@ -30,13 +32,14 @@ def load_model_tokenizer():
generation_config = GenerationConfig.from_pretrained(model_id) generation_config = GenerationConfig.from_pretrained(model_id)
return model, tokenizer, generation_config return model, tokenizer, generation_config
def clear_chat_messages(): def clear_chat_messages():
del st.session_state.messages del st.session_state.messages
def init_chat_messages(): def init_chat_messages():
with st.chat_message("assistant", avatar='🤖'): with st.chat_message("assistant", avatar='🤖'):
st.markdown("您好我是由Joya开发的MiniMind很高兴为您服务😄") st.markdown("您好我是由JingyaoGong创造的MiniMind很高兴为您服务😄")
if "messages" in st.session_state: if "messages" in st.session_state:
for message in st.session_state.messages: for message in st.session_state.messages:
@ -48,6 +51,7 @@ def init_chat_messages():
return st.session_state.messages return st.session_state.messages
# max_new_tokens = st.sidebar.slider("max_new_tokens", 0, 1024, 512, step=1) # max_new_tokens = st.sidebar.slider("max_new_tokens", 0, 1024, 512, step=1)
# top_p = st.sidebar.slider("top_p", 0.0, 1.0, 0.8, step=0.01) # top_p = st.sidebar.slider("top_p", 0.0, 1.0, 0.8, step=0.01)
# top_k = st.sidebar.slider("top_k", 0, 100, 0, step=1) # top_k = st.sidebar.slider("top_k", 0, 100, 0, step=1)
@ -67,7 +71,7 @@ def main():
placeholder = st.empty() placeholder = st.empty()
chat_messages = [] chat_messages = []
chat_messages.append({"role": "user", "content": prompt}) chat_messages.append({"role": "user", "content": '请问,' + prompt})
# print(messages) # print(messages)
new_prompt = tokenizer.apply_chat_template( new_prompt = tokenizer.apply_chat_template(
chat_messages, chat_messages,