saenoooos/BlockMasterAI-Playground

1import time2from datetime import datetime3import os4import streamlit as st5import torch6from transformers import AutoModelForCausalLM, AutoTokenizer7import mysql.connector8import re9 10# Initialize CUDA11device = torch.device('cuda')12 13# Model and tokenizer14name = "saenoooos/BlockMasterAI-Chatbot"15 16 17@st.cache_resource18def load_model():19    mistral_model = AutoModelForCausalLM.from_pretrained(20        name,21        device_map='cuda'22    )23    mistral_tokenizer = AutoTokenizer.from_pretrained(name)24    return mistral_model, mistral_tokenizer25 26 27model, tokenizer = load_model()28 29# Database Info30dbname = os.environ["dbname"]31dbhost = os.environ["dbhost"]32dbuser = os.environ["dbuser"]33dbpass = os.environ["dbpass"]34 35 36def database_login():37    database = mysql.connector.connect(38        host=dbhost,39        user=dbuser,40        passwd=dbpass,41    )42    return database43 44 45# Fetch previous messages from the database46def fetch_previous_messages(token):47    database = database_login()48    cursor = database.cursor()49    messages = []50    if re.match(r'^[\w-]+$', token):51        try:52            cursor.execute(f"USE `{dbname}`;")53            query = f"SELECT * FROM `token_{token}`;"54            cursor.execute(query)55            results = cursor.fetchall()56            for result in results:57                messages.append({"role": result[1], "content": result[2]})58            return messages59        except mysql.connector.Error as err:60            st.error(f"Database Error: {err}")61    else:62        st.error("Invalid token in fetch message")63    return messages64 65 66def push_message(msgtype, message, token):67    database = database_login()68    cursor = database.cursor()69    if re.match(r'^[\w-]+$', token):70        try:71            cursor.execute(f"USE `{dbname}`;")72            sql_query = f"INSERT INTO token_{token} (role, content,created_at) VALUES (%s,%s,%s)"73            val = (msgtype, message, datetime.now())74 75            cursor.execute(sql_query, val)76            database.commit()77        except mysql.connector.Error as err:78            st.error(f"Database Error: {err}")79    else:80        st.error("Invalid token in push message")81 82 83def format_message_history(token):84    database = database_login()85    cursor = database.cursor()86    messages = []87    if re.match(r'^[\w-]+$', token):88        try:89            cursor.execute(f"USE `{dbname}`;")90            query = f"SELECT * FROM `token_{token}`;"91            cursor.execute(query)92            results = cursor.fetchall()93            for result in results:94                messages.append({"role": result[1], "content": result[2]})95            return messages96        except mysql.connector.Error as err:97            st.error(f"Database Error: {err}")98    else:99        st.error("Invalid token in fetch message")100    return messages101 102 103# Generate response using the model104def mistral_inference(prompt, token):105    with st.spinner("Accessing database..."):106        messages = fetch_previous_messages(token)107        if messages is None:108            messages = []109        # Append the user's input to the conversation history110        messages.append({"role": "user", "content": prompt})111 112    # Tokenize the conversation history once113    with st.spinner("Applying tokenization template to chat history..."):114        with torch.autocast(device_type='cuda', dtype=torch.float16):115            model_inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to("cuda")116 117    with st.spinner("Generating new tokens..."):118        # Generate the assistant's response with efficient sampling119        generated_ids = model.generate(120            model_inputs,121            max_new_tokens=500,122            do_sample=user_do_sample,123            top_k=user_top_k,  # top_k sampling for efficiency124            top_p=user_top_p,  # Nucleus sampling125            temperature=user_temperature,  # Control creativity126        )127    with st.spinner("Batch decoding new tokens..."):128        new_tokens = generated_ids[:, model_inputs.shape[-1]:]129        result = tokenizer.batch_decode(new_tokens, skip_special_tokens=True)130 131    with st.spinner("Pushing result to database..."):132        # Push user prompt & assistant response to DB133        push_message("user", prompt, token)134        push_message("assistant", result[0], token)135 136    with st.spinner("Returning result..."):137        return result[0]138 139 140# Streamlit interface141 142user_token = "testing"143 144st.markdown("### Welcome to the BlockMasterAI testing playground!")145chatbox = st.container(border=True,height=450)146with chatbox:147    with st.spinner("Fetching history..."):148        current = format_message_history(user_token)149        for msg in current:150            if msg["role"] == "system":151                systembox = st.chat_message("system", avatar="blockmaster-chat-system.png")152                systembox.write(msg["content"])153            if msg["role"] == "user":154                userbox = st.chat_message("user", avatar="blockmaster-chat-user.png")155                userbox.write(msg["content"])156            if msg["role"] == "assistant":157                assistantbox = st.chat_message("assistant", avatar="blockmaster-chat-assistant.png")158                assistantbox.write(msg["content"])159 160        # st.image("blockmaster-chat-assistant.png", width=100)161        # Input Section162user_prompt = st.chat_input("Enter your prompt:")163 164 165st.markdown("""166 167#### Please read the following carefully!168This webpage allows you to directly test the BlockMasterAI integrated chatbot without having to run it inside of Minecraft! It offers a user-friendly design enabling seamless interaction between you (the user), and the bot (BlockMasterAI).169 170Please keep in mind that every message sent here will be saved and added to the same chat history! This is by design, as it allows me to collect the data on responses to allow for greater tuning.171 172With that in mind, **do not input any personal identifying information - passwords, phone numbers, etc into this interface as it will be very difficult to manually remove!**173I take no responsibility for the messages sent to this interface!174 175All messages sent are saved, but we don't want to be displaying thousands here! Please note that refreshing the page will cause you to lose any messages sent from your local history.176 177In the instance that you need assistance with removing any information please drop me an email with the deets and I'll try my best.178 179Have fun exploring the capabilities of large language models!180""")181st.html("Stumbled here by accident? Return to the <a href='http://blockmaster.minecraftengineering.org'>main page.</a><br/>")182 183st.image("https://saen.minecraftengineering.org/projects/blockmaster/sheeppunch.png",width=400)184 185with st.sidebar:186    st.markdown("#### Debug Parameters")187    st.write("Feel free to play around with adjusting these values to see how the output changes! If you don't know what any of this means then don't worry, you don't have to change anything here to get the model to work!")188 189    st.write(f"Running on {torch.cuda.get_device_name(torch.cuda.current_device())}. CUDA Available: {torch.cuda.is_available()}.")190 191    user_top_k = st.slider("Top-K Sampling (default 100)",value=100,min_value=1,max_value=100)192    st.info("Restricts the sampling space to the K most likely next words. This narrows down the potential candidates and reduces the chance of selecting improbable words.")193    user_top_p = st.slider("Top-P Nucleus Sampling (default 0.99)",value=0.99,min_value=0.01,max_value=1.00,step=0.01)194    st.info("Selects the smallest set of words whose cumulative probability exceeds the threshold P. This is a more dynamic approach than Top-K.")195    user_do_sample = st.toggle("do_sample (default True)",value=True)196    st.info("Sets whether the model uses sampling or greedy decoding during text generation. When do_sample is set to True, the model samples the next word from a probability distribution, introducing randomness and creativity. When do_sample is False, the model selects the most probable word at each step, resulting in more deterministic and predictable outputs.")197    user_temperature = st.slider("Temperature (default 1.00)",value=1.0,min_value=0.1,max_value=1.00,step=0.1)198    st.info("Controls the randomness of the words the model chooses. Lower temperatures produce less creative responses, while higher temperatures produce more diverse or creative results.")199 200 201 202def convert_response_streamable(strinput):203    for word in strinput.split(" "):204        yield word + " "205        time.sleep(0.02)206 207 208if user_prompt:209    with chatbox:210        st.chat_message("user",avatar="blockmaster-chat-user.png").write(user_prompt)211        with st.spinner("Generating response (this may take up to 60 seconds) ..."):212            try:213                response = mistral_inference(user_prompt, user_token)214                success = st.success("Response generated successfully!")215                st.chat_message("assistant",avatar="blockmaster-chat-assistant.png").write_stream(convert_response_streamable(response))216                time.sleep(3)217                success.empty()218            except Exception as e:219                st.error(f"An error occurred: {e}")220