Spaces:
Runtime error
Runtime error
File size: 1,665 Bytes
2907c25 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import os
from llama_cpp_chat_completion_wrapper import Llama2ChatCompletionWrapper, Message
import streamlit as st
def console_print(message: Message) -> None:
reset = "\033[00m"
color_map = {
"system": ("\033[1;35m", "\033[35m"),
"user": ("\033[1;33m", "\033[33m"),
"assistant": ("\033[1;31m", "\033[31m"),
"assistant-before-post-process": ("\033[1;31m", "\033[31m"),
}
role_color, content_color = color_map[message["role"]]
formatted_message = f"{role_color}{message['role'].upper()}{reset}> {content_color}{message['content']}{reset}"
print(formatted_message)
@st.cache_resource
def load_model():
model_path = os.path.join(os.path.dirname(__file__), "llama-2-7b-chat.ggmlv3.q2_K.bin")
return Llama2ChatCompletionWrapper(model_path=model_path, callback=console_print)
def main():
if len(st.session_state) == 0:
st.session_state["llm"] = load_model()
st.session_state["params"] = {
"temp": 0,
"top_p": 1,
"frequency_penalty": 0,
"presence_penalty": 0,
}
st.session_state["llm"].new_session()
st.session_state["messages"] = st.session_state["llm"].messages
st.title("💬 Chatbot")
for msg in st.session_state["messages"]:
st.chat_message(msg["role"]).write(msg["content"])
if prompt := st.chat_input():
st.chat_message("user").write(prompt)
response = st.session_state["llm"](prompt, params=st.session_state["params"])
st.chat_message("assistant").write(response)
if __name__ == "__main__":
main()
|