# adapted from: # https://medium.com/@james.irving.phd/creating-your-personal-chatbot-using-hugging-face-spaces-and-streamlit-596a54b9e3ed import os from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, pipeline from transformers import LlamaTokenizer import streamlit as st import torch # Define the model repository REPO_NAME = 'schuler/experimental-JP47D20' # REPO_NAME = 'schuler/experimental-JP47D21-KPhi-3-micro-4k-instruct' # Configure the Streamlit app st.set_page_config(page_title="Experimental KPhi3 Model - Currently in Training", page_icon="🤗") st.title("Experimental KPhi3 Model - Currently in Training") # Load tokenizer and model @st.cache_resource(show_spinner="Loading model...") def load_model(local_repo_name): # tokenizer = AutoTokenizer.from_pretrained(local_repo_name, trust_remote_code=True) tokenizer = LlamaTokenizer.from_pretrained(local_repo_name, trust_remote_code=True) generator_conf = GenerationConfig.from_pretrained(local_repo_name) model = AutoModelForCausalLM.from_pretrained(local_repo_name, trust_remote_code=True, torch_dtype=torch.bfloat16) return tokenizer, generator_conf, model tokenizer, generator_conf, model = load_model(REPO_NAME) total_params = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) embed_params = sum(p.numel() for p in model.model.embed_tokens.parameters())*2 non_embed_params = (trainable_params - embed_params) / 1e6 st.markdown(f"*This chat uses the {REPO_NAME} model with {model.get_memory_footprint() / 1e6:.2f} MB memory footprint. ") # st.markdown(f"Total number of parameters: {total_params}. ") # st.markdown(f"Total number of trainable parameters: {trainable_params}. ") # st.markdown(f"Total number of embed parameters: {embed_params}. ") st.markdown(f"Total number of non embedding trainable parameters: {non_embed_params:.2f} million. ") st.markdown(f"You may ask questions such as 'What is biology?' or 'What is the human body?'*") try: generator = pipeline("text-generation", model=model, tokenizer=tokenizer) except Exception as e: st.error(f"Failed to load model: {str(e)}") # 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 = 64 if "system_message" not in st.session_state: st.session_state.system_message = "" if "starter_message" not in st.session_state: st.session_state.starter_message = "Hello, there! How can I help you today?" if "can_continue" not in st.session_state: st.session_state.can_continue = False # Initialize state for continue action need_continue = False # Initialize the last response if "last_response" not in st.session_state: st.session_state.last_response = '' # Sidebar for settings with st.sidebar: st.header("System Settings") # AI Settings st.session_state.system_message = st.text_area( "System Message", value=st.session_state.system_message ) st.session_state.starter_message = st.text_area( 'First AI Message', value=st.session_state.starter_message ) # Model Settings st.session_state.max_response_length = st.number_input( "Max Response Length", value=st.session_state.max_response_length ) # 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, max_new_tokens=256, continue_last=False): """ 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. max_new_tokens (int): The maximum number of new tokens to generate. continue_last (bool): Whether to continue the last assistant response. Returns: tuple: A tuple containing the generated response and the updated chat history. """ if continue_last: # We want to continue the last assistant response prompt = st.session_state.last_response else: # Build the conversation prompt if (len(system_message)>0): prompt = "<|assistant|>"+system_message+f"<|end|>" else: prompt = '' # f"{system_message}\nCurrent Conversation:\n" for message in chat_history: role = "<|assistant|>" if message['role'] == 'assistant' else "<|user|>" prompt += f"{role}{message['content']}<|end|>" prompt += f"<|user|>{user_text}<|end|><|assistant|>" # Generate the response response_output = generator( prompt, generation_config=generator_conf, max_new_tokens=max_new_tokens, do_sample=True, top_p=0.25, repetition_penalty=1.2 ) generated_text = response_output[0]['generated_text'] st.session_state.last_response = generated_text # Extract the assistant's response assistant_response = generated_text[len(prompt):] # .strip() if continue_last: # Append the continued text to the last assistant message st.session_state.chat_history[-1]['content'] += assistant_response else: # Update the chat history chat_history.append({'role': 'user', 'content': user_text}) chat_history.append({'role': 'assistant', 'content': assistant_response}) return assistant_response, chat_history # Chat interface chat_interface = st.container() def refresh_chat(): with chat_interface: output_container = st.container() # Display chat messages with output_container: for idx, message in enumerate(st.session_state.chat_history): if message['role'] == 'system': continue with st.chat_message(message['role'], avatar=st.session_state.avatars[message['role']]): st.markdown(message['content']) # If this is the last assistant message, add the "Continue" button # if idx == len(st.session_state.chat_history) - 1 and message['role'] == 'assistant': refresh_chat() # User input area (moved to the bottom) st.session_state.user_text = st.chat_input(placeholder="Enter your text here.") # When the user enters new text if st.session_state.user_text: # Display the user's message with st.chat_message("user", avatar=st.session_state.avatars['user']): st.markdown(st.session_state.user_text) # Display a spinner while generating the response with st.chat_message("assistant", avatar=st.session_state.avatars['assistant']): with st.spinner("Thinking..."): # Generate the assistant's response 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, continue_last=False ) st.markdown(response) st.session_state.can_continue = True # Clear the user input st.session_state.user_text = None if st.session_state.can_continue: if st.button("Continue"): need_continue = True else: need_continue = False # If "Continue" button was pressed if need_continue: # Display a spinner while generating the continuation with st.chat_message("assistant", avatar=st.session_state.avatars['assistant']): with st.spinner("Continuing..."): # Generate the continuation of the assistant's last response response, st.session_state.chat_history = get_response( system_message=st.session_state.system_message, user_text=None, chat_history=st.session_state.chat_history, max_new_tokens=st.session_state.max_response_length, continue_last=True ) st.markdown(response) st.rerun()