import streamlit as st import os from streamlit_extras.stylable_container import stylable_container from PIL import Image from langchain_huggingface import HuggingFaceEndpoint from langchain_core.prompts import PromptTemplate from langchain_core.output_parsers import StrOutputParser # Variables used Globally path = "/data" #preset path for hugging face spaces for persistent storage and cannot be changed model_id="mistralai/Mistral-7B-Instruct-v0.3" # Configure the Streamlit app st.set_page_config(page_title="CRIStine", page_icon = "👩‍🔧") st.title("CRIStine - Interactive CRIS Assistant") st.markdown(f"*This is a chatbot that uses the HuggingFace transformers with Retrieval Augmented Generation to guide and train users. It uses the {model_id}.*") # Application Functions # File Loader @st.dialog("Upload a File") def upload_file(): uploaded_file = st.file_uploader("Choose a file") if uploaded_file is not None: file_details = {"FileName":uploaded_file.name,"FileType":uploaded_file.type} st.write(file_details) with open(os.path.join(path,uploaded_file.name),"wb") as f: f.write(uploaded_file.getbuffer()) st.success("Saved File") # File Delete @st.dialog("Delete a File") def delete_file(): # List all files in directory and subdirectories as buttons for root, dirs, file_names in os.walk(path): for file_name in file_names: if st.button(file_name): os.remove(os.path.join(path,file_name)) st.success("Removed File") st.rerun() # File View @st.dialog("Files used by AI") def view_file(): # List all files in directory and subdirectories files = [] for root, dirs, file_names in os.walk(path): for file_name in file_names: files.append(file_name) st.write(files) if st.button("Close"): st.rerun() logo_column, space_column, upload_column, delete_column, browse_column, recycle_column = st.columns(6) st.markdown( ' ', unsafe_allow_html=True, ) with logo_column: image = Image.open(os.path.join(path,'CRIStine.png')) st.image(image, caption='CRIStine') with upload_column: with stylable_container( key="upload_button", css_styles=r""" button p:before { font-family: 'Font Awesome 5 Free'; content: '\f574'; display: inline-block; padding-right: 3px; vertical-align: middle; font-weight: 900; } """, ): if st.button("Upload", key='upload'): upload_file() with delete_column: with stylable_container( key="delete_button", css_styles=r""" button p:before { font-family: 'Font Awesome 5 Free'; content: '\f1c3'; display: inline-block; padding-right: 3px; vertical-align: middle; font-weight: 900; } """, ): if st.button("Delete", key='delete'): delete_file() with browse_column: with stylable_container( key="view_button", css_styles=r""" button p:before { font-family: 'Font Awesome 5 Free'; content: '\f07c'; display: inline-block; padding-right: 3px; vertical-align: middle; font-weight: 900; } """, ): if st.button("View", key='view'): view_file() with recycle_column: with stylable_container( key="recycle_button", css_styles=r""" button p:before { font-family: 'Font Awesome 5 Free'; content: '\f1b8'; display: inline-block; padding-right: 3px; vertical-align: middle; font-weight: 900; } """, ): st.button("Recycle", key='recycle') # Main app goes below here - def get_llm_hf_inference(model_id=model_id, max_new_tokens=128, temperature=0.1): """ Returns a language model for HuggingFace inference. Parameters: - model_id (str): The ID of the HuggingFace model repository. - max_new_tokens (int): The maximum number of new tokens to generate. - temperature (float): The temperature for sampling from the model. Returns: - llm (HuggingFaceEndpoint): The language model for HuggingFace inference. """ llm = HuggingFaceEndpoint( repo_id=model_id, max_new_tokens=max_new_tokens, temperature=temperature, token = os.getenv("HF_TOKEN") ) return llm # Initialize session state for avatars if "avatars" not in st.session_state: st.session_state.avatars = {'user': None, 'assistant': None} # Initialize session state for user text input if 'user_text' not in st.session_state: st.session_state.user_text = None # Initialize session state for model parameters if "max_response_length" not in st.session_state: st.session_state.max_response_length = 256 if "system_message" not in st.session_state: st.session_state.system_message = "friendly AI conversing with a human user" if "starter_message" not in st.session_state: st.session_state.starter_message = "Hello, there! How can I help you today?" # Sidebar for settings with st.sidebar: st.header("System Settings") # AI Settings st.session_state.system_message = st.text_area( "System Message", value="You are a friendly AI conversing with a human user." ) st.session_state.starter_message = st.text_area( 'First AI Message', value="Hello, there! How can I help you today?" ) # Model Settings st.session_state.max_response_length = st.number_input( "Max Response Length", value=128 ) # Avatar Selection st.markdown("*Select Avatars:*") col1, col2 = st.columns(2) with col1: st.session_state.avatars['assistant'] = st.selectbox( "AI Avatar", options=["🤗", "💬", "🤖"], index=0 ) with col2: st.session_state.avatars['user'] = st.selectbox( "User Avatar", options=["👤", "👱‍♂️", "👨🏾", "👩", "👧🏾"], index=0 ) # Reset Chat History reset_history = st.button("Reset Chat History") # Initialize or reset chat history if "chat_history" not in st.session_state or reset_history: st.session_state.chat_history = [{"role": "assistant", "content": st.session_state.starter_message}] def get_response(system_message, chat_history, user_text, eos_token_id=['User'], max_new_tokens=256, get_llm_hf_kws={}): """ Generates a response from the chatbot model. Args: system_message (str): The system message for the conversation. chat_history (list): The list of previous chat messages. user_text (str): The user's input text. model_id (str, optional): The ID of the HuggingFace model to use. eos_token_id (list, optional): The list of end-of-sentence token IDs. max_new_tokens (int, optional): The maximum number of new tokens to generate. get_llm_hf_kws (dict, optional): Additional keyword arguments for the get_llm_hf function. Returns: tuple: A tuple containing the generated response and the updated chat history. """ # Set up the model hf = get_llm_hf_inference(max_new_tokens=max_new_tokens, temperature=0.1) # Create the prompt template prompt = PromptTemplate.from_template( ( "[INST] {system_message}" "\nCurrent Conversation:\n{chat_history}\n\n" "\nUser: {user_text}.\n [/INST]" "\nAI:" ) ) # Make the chain and bind the prompt chat = prompt | hf.bind(skip_prompt=True) | StrOutputParser(output_key='content') # Generate the response response = chat.invoke(input=dict(system_message=system_message, user_text=user_text, chat_history=chat_history)) response = response.split("AI:")[-1] # Update the chat history chat_history.append({'role': 'user', 'content': user_text}) chat_history.append({'role': 'assistant', 'content': response}) return response, chat_history # Chat interface chat_interface = st.container(border=True) with chat_interface: output_container = st.container() st.session_state.user_text = st.chat_input(placeholder="Enter your text here.") # Display chat messages with output_container: # For every message in the history for message in st.session_state.chat_history: # Skip the system message if message['role'] == 'system': continue # Display the chat message using the correct avatar with st.chat_message(message['role'], avatar=st.session_state['avatars'][message['role']]): st.markdown(message['content']) # When the user enter new text: if st.session_state.user_text: # Display the user's new message immediately with st.chat_message("user", avatar=st.session_state.avatars['user']): st.markdown(st.session_state.user_text) # Display a spinner status bar while waiting for the response with st.chat_message("assistant", avatar=st.session_state.avatars['assistant']): with st.spinner("Thinking..."): # Call the Inference API with the system_prompt, user text, and history response, st.session_state.chat_history = get_response( system_message=st.session_state.system_message, user_text=st.session_state.user_text, chat_history=st.session_state.chat_history, max_new_tokens=st.session_state.max_response_length, ) st.markdown(response)