Update pages/1_Simple_Chat_UI.py

pages/1_Simple_Chat_UI.py

@@ -2,12 +2,12 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
 
 def load_model_tokenizer(model_name, hf_api_key):
-    if model_name == "Mistral-7B":
-        model_name="NousResearch/Nous-Hermes-2-Mistral-7B-DPO"
+    if model_name == "LLaMa-2B":
+        model_name="llmware/bling-sheared-llama-2.7b-0.1"
         model = AutoModelForCausalLM.from_pretrained(model_name, token=hf_api_key)
         tokenizer = AutoTokenizer.from_pretrained(model_name, tokenizer=hf_api_key)
-    elif model_name == "blenderbot-400M-distill":
-        model_name = "facebook/blenderbot-400M-distill"
+    elif model_name == "Red-Pajamas-3b":
+        model_name = "llmware/bling-red-pajamas-3b-0.1"
         model = AutoModelForCausalLM.from_pretrained(model_name)
         tokenizer = AutoTokenizer.from_pretrained(model_name)
     return (model,tokenizer)
@@ -38,7 +38,7 @@ with st.sidebar:
 if "messages" not in st.session_state.keys():
     st.session_state.messages = [{"role": "assistant", "content": "How may I help you?"}]
 
-model_name = st.radio("Select model to chat", options=["Mistral-7B", "LLaMa-2B", "blenderbot-400M-distill"], horizontal=True, key='model_selection')
+model_name = st.radio("Select model to chat", options=["LLaMa-2B", "Red-Pajamas-3b"], horizontal=True, key='model_selection')
 model, tokenizer = load_model_tokenizer(model_name, hf_api_key)
 
 for message in st.session_state.messages: