serafm/chatbot

1import streamlit as st2from llama_index import(3    SimpleDirectoryReader,4    VectorStoreIndex,5    SeviceContext,6)7from llama_index.llms import LlamaCPP8from llama_index.llms.llama_utils import(9    messages_to_prompt,10    completion_to_prompt,11)12from langchain.schema import(HumanMessage, AIMessage)13 14def init_page() -> None:15    st.set_page_config(16        page_title = "ChatTopG"17    )18    st.header("ChatTopG")19    st.sidebar.title("Options")20 21def select_llm() -> LlamaCPP:22    return LlamaCPP(23        model_path = "/home/serafm/Documents/GitHub/hugging-face-chatbot/llama-2-7b-chat.Q2_K.gguf",24        temperature = 0.1,25        max_new_token = 500,26        context_window = 3900,27        generate_kwargs = {},28        model_kwargs = {"n_gpu_layers":1},29        messages_to_prompt = messages_to_prompt,30        completion_to_prompt = completion_to_prompt,31        verbose = True,32    )33 34def init_messages() -> None:35    clear_button = st.sidebar.button("Clear Conversation", key="clear")36    if clear_button or "messages" not in st.session_state:37        st.session_state.messages = [38            SystemMessage(39                content = "You are a helpful AI assistant. Reply your answer in markdown format."40            )41        ]42 43def get_answer(llm, messages) -> str:44    response = llm.complete(messages)45    return response.text46 47def main() -> None:48    init_page()49    llm = select_llm()50    init_messages()51 52    if user_input := st.chat_input("Input your question"):53        st.session_state.messages.append(HumanMessage(content=user_input))54        with st.spinner("Bot is typing..."):55            answer = get_answer(llm, user_input)56            print(answer)57        st.session_state.messages.append(AIMessage(content=answer))58 59    messages = st.session_state.get("messages", [])60    for message in messages:61        if isinstance(message, AIMessage):62            with st.chat_message("assistant"):63                st.markdown(message.content)64        elif isinstance(message, HumanMessage):65            with st.chat_message("user"):66                st.markdown(message.content)67 68 69 70if __name__ == "__main__":71    main()