Update pages/1_Simple_Chat_UI.py

pages/1_Simple_Chat_UI.py

@@ -1,60 +1,59 @@
-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="mistralai/Mistral-7B-Instruct-v0.2"
-        #model = AutoModelForCausalLM.from_pretrained(model_name, token=hf_api_key)
-        #tokenizer = AutoTokenizer.from_pretrained(model_name, tokenizer=hf_api_key)
-        pass
-    elif model_name == "blenderbot-400M-distill":
-        model_name = "facebook/blenderbot-400M-distill"
-        model = AutoModelForCausalLM.from_pretrained(model_name)
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-    return (model,tokenizer)
-
-def generate_response(prompt_input, model, tokenizer):
-    inputs = tokenizer.encode_plus(prompt_input, return_tensors="pt")
-    # Generate the response from the model with additional parameters
-    outputs = model.generate(**inputs, max_length=max_length, do_sample=True ,temperature=temperature)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
-    return response
-
-st.set_page_config(page_title="Learn Geoscience")
-
-with st.sidebar:
-    st.title('Learn Geoscience Chat')
-    if 'hf_key' in st.secrets:
-        st.success('Huggingface API key provided', icon='✅')
-        hf_api_key = st.secrets['hf_key']
-    else:
-        hf_api_key = st.text_input('Enter Huggingface API Key:', type='password')
-        if not hf_api_key:
-            st.warning('Please enter Huggingface API key!', icon='⚠️')
-        else:
-            st.success('Proceed to entering your prompt message!', icon='👉')
-    max_length = st.slider("Max Length", 10, 100, 50)
-    temperature = st.slider("Temperature", 0.0, 1.0, 0.7)
-
-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, tokenizer = load_model_tokenizer(model_name, hf_api_key)
-
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.write(message["content"])
-
-if prompt := st.chat_input(disabled = not hf_api_key):
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    with st.chat_message("user"):
-        st.write(prompt)
-
-if st.session_state.messages[-1]["role"] != "assistant":
-    with st.chat_message("assistant"):
-        with st.spinner("Thinking..."):
-            response = generate_response(prompt, model, tokenizer)
-            st.write(response)
-    message = {"role": "assistant", "content": response}
-    st.session_state.messages.append(message)
+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="mistralai/Mistral-7B-Instruct-v0.2"
+        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"
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+    return (model,tokenizer)
+
+def generate_response(prompt_input, model, tokenizer):
+    inputs = tokenizer.encode_plus(prompt_input, return_tensors="pt")
+    # Generate the response from the model with additional parameters
+    outputs = model.generate(**inputs, max_length=max_length, do_sample=True ,temperature=temperature)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+    return response
+
+st.set_page_config(page_title="Learn Geoscience")
+
+with st.sidebar:
+    st.title('Learn Geoscience Chat')
+    if 'hf_key' in st.secrets:
+        st.success('Huggingface API key provided', icon='✅')
+        hf_api_key = st.secrets['hf_key']
+    else:
+        hf_api_key = st.text_input('Enter Huggingface API Key:', type='password')
+        if not hf_api_key:
+            st.warning('Please enter Huggingface API key!', icon='⚠️')
+        else:
+            st.success('Proceed to entering your prompt message!', icon='👉')
+    max_length = st.slider("Max Length", 10, 100, 50)
+    temperature = st.slider("Temperature", 0.0, 1.0, 0.7)
+
+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, tokenizer = load_model_tokenizer(model_name, hf_api_key)
+
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.write(message["content"])
+
+if prompt := st.chat_input(disabled = not hf_api_key):
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.write(prompt)
+
+if st.session_state.messages[-1]["role"] != "assistant":
+    with st.chat_message("assistant"):
+        with st.spinner("Thinking..."):
+            response = generate_response(prompt, model, tokenizer)
+            st.write(response)
+    message = {"role": "assistant", "content": response}
+    st.session_state.messages.append(message)