Update app.py

app.py

@@ -1,37 +1,15 @@
-### app.py code is taken from https://huggingface.co/spaces/ngebodh/SimpleChatbot-Backup/blob/main/app.py
-### https://medium.com/@nigelgebodh/large-language-models-chatting-with-ai-chatbots-from-google-mistral-ai-and-hugging-face-b33efedea38d
-
-""" Simple Chatbot
-@author: Sagar Padhiyar
-@email: spadhiyar230595@gmail.com
-"""
-
 import streamlit as st
-from openai import OpenAI
+from huggingface_hub import InferenceClient
 import os
 import sys
-from dotenv import load_dotenv, dotenv_values
-from huggingface_hub import InferenceClient
-load_dotenv()
-
-
-
-
-
 
-# # initialize the client
-# client = OpenAI(
-#   base_url="https://api-inference.huggingface.co/v1",
-#   api_key=os.environ.get('HUGGINGFACE_API')#"hf_xxx" # Replace with your token
-# ) 
+st.title("ChatGPT-like Chat boat")
 
-
-base_url="https://api-inference.huggingface.co/v1"
+base_url="https://api-inference.huggingface.co/models/"
 
 API_KEY = os.environ.get('HUGGINGFACE_API')
+headers = {"Authorization":"Bearer "+API_KEY}
 
-
-#Create supported models
 model_links ={
     "Mistral":base_url+"mistralai/Mistral-7B-Instruct-v0.2",
     "Gemma-7B":base_url+"google/gemma-7b-it",
@@ -40,8 +18,6 @@ model_links ={
     # "Llama-2":"meta-llama/Llama-2-7b-chat-hf"
 }
 
-headers = {"Authorization":"Bearer "+API_KEY}
-
 #Pull info about the model to display
 model_info ={
     "Mistral":
@@ -75,6 +51,13 @@ model_info ={
 
 }
 
+def format_promt(message, custom_instructions=None):
+    prompt = ""
+    if custom_instructions:
+        prompt += f"[INST] {custom_instructions} [/INST]"
+    prompt += f"[INST] {message} [/INST]"
+    return prompt
+
 def reset_conversation():
     '''
     Resets Conversation
@@ -82,16 +65,7 @@ def reset_conversation():
     st.session_state.conversation = []
     st.session_state.messages = []
     return None
-    
-def format_promt(message, custom_instructions=None):
-    if custom_instructions:
-        promt += f"[INST] {custom_instructions} [/INST]"
-    promt += f"[INST] {message} [/INST]"
-    return promt
-
-
 
-# Define the available models
 models =[key for key in model_links.keys()]
 
 # Create the sidebar with the dropdown for model selection
@@ -100,11 +74,9 @@ selected_model = st.sidebar.selectbox("Select Model", models)
 #Create a temperature slider
 temp_values = st.sidebar.slider('Select a temperature value', 0.0, 1.0, (0.5))
 
-
 #Add reset button to clear conversation
 st.sidebar.button('Reset Chat', on_click=reset_conversation) #Reset button
 
-
 # Create model description
 st.sidebar.write(f"You're now chatting with **{selected_model}**")
 st.sidebar.markdown(model_info[selected_model]['description'])
@@ -112,8 +84,6 @@ st.sidebar.image(model_info[selected_model]['logo'])
 st.sidebar.markdown("*Generated content may be inaccurate or false.*")
 st.sidebar.markdown("\nLearn how to build this chatbot [here](https://ngebodh.github.io/projects/2024-03-05/).")
 
-
-
 if "prev_option" not in st.session_state:
     st.session_state.prev_option = selected_model
 
@@ -123,44 +93,34 @@ if st.session_state.prev_option != selected_model:
     st.session_state.prev_option = selected_model
     reset_conversation()
 
-
-
 #Pull in the model we want to use
 repo_id = model_links[selected_model]
 
-
 st.subheader(f'AI - {selected_model}')
 # st.title(f'ChatBot Using {selected_model}')
 
-# Set a default model
-if selected_model not in st.session_state:
-    st.session_state[selected_model] = model_links[selected_model] 
-
 # Initialize chat history
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-
 # Display chat messages from history on app rerun
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
 
 
-
 # Accept user input
 if prompt := st.chat_input(f"Hi I'm {selected_model}, ask me a question"):
 
     custom_instruction = "Act like a Human in conversation"
 
-    formated_text = format_promt(prompt, custom_instruction)
-
     # Display user message in chat message container
     with st.chat_message("user"):
         st.markdown(prompt)
     # Add user message to chat history
     st.session_state.messages.append({"role": "user", "content": prompt})
 
+    formated_text = format_promt(prompt, custom_instruction)
 
     # Display assistant response in chat message container
     with st.chat_message("assistant"):
@@ -171,8 +131,10 @@ if prompt := st.chat_input(f"Hi I'm {selected_model}, ask me a question"):
         output = client.text_generation(
             formated_text,
             temperature=temp_values,#0.5
-            max_tokens=3000,
+            max_new_tokens=3000,
+            stream=True
         )
 
         response = st.write_stream(output)
-    st.session_state.messages.append({"role": "assistant", "content": response})
+    st.session_state.messages.append({"role": "assistant", "content": response})
+