Update

Dockerfile

@@ -0,0 +1,34 @@
+# read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
+# you will also find guides on how best to write your Dockerfile
+
+FROM nvidia/cuda:12.2.0-devel-ubuntu20.04
+
+#set up environment
+RUN apt-get update && apt-get install --no-install-recommends --no-install-suggests -y curl
+RUN apt-get install unzip
+RUN apt-get -y install python3
+RUN apt-get -y install python3-pip
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN pip3 install --no-cache-dir --upgrade -r /code/requirements.txt
+
+# Set up a new user named "user" with user ID 1000
+RUN useradd -m -u 1000 user
+
+# Switch to the "user" user
+USER user
+
+# Set home to the user's home directory
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+# Set the working directory to the user's home directory
+WORKDIR $HOME/app
+
+# Copy the current directory contents into the container at $HOME/app setting the owner to the user
+COPY --chown=user . $HOME/app
+
+CMD ["streamlit", "run", "app.py", "--server.port", "7860", "--server.address", "0.0.0.0"]

app.py

@@ -0,0 +1,71 @@
+import json
+import torch
+import streamlit as st
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers.generation.utils import GenerationConfig
+
+
+st.set_page_config(page_title="Baichuan-13B-Chat")
+st.title("Baichuan-13B-Chat")
+
+@st.cache_resource
+def init_model():
+    model = AutoModelForCausalLM.from_pretrained(
+        "baichuan-inc/Baichuan-13B-Chat",
+        torch_dtype=torch.float16,
+        device_map="auto",
+        trust_remote_code=True
+    )
+    model.generation_config = GenerationConfig.from_pretrained(
+        "baichuan-inc/Baichuan-13B-Chat"
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        "baichuan-inc/Baichuan-13B-Chat",
+        use_fast=False,
+        trust_remote_code=True
+    )
+    return model, tokenizer
+
+
+def clear_chat_history():
+    del st.session_state.messages
+
+
+def init_chat_history():
+    with st.chat_message("assistant", avatar='🤖'):
+        st.markdown("Greetings! I am the BaiChuan large language model, delighted to assist you.🥰")
+
+    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
+
+
+def main():
+    model, tokenizer = init_model()
+    messages = init_chat_history()
+
+    if prompt := st.chat_input("Shift + Enter for a new line, Enter to send"):
+        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()
+            for response in model.chat(tokenizer, messages, stream=True):
+                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("Reset Chat", on_click=clear_chat_history)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,6 @@
+accelerate
+colorama
+cpm_kernels
+sentencepiece
+streamlit
+transformers_stream_generator