Update app.py

app.py

@@ -1,8 +1,11 @@
 import streamlit as st
 import os
+import torch
 from openai import OpenAI
-import random
+import numpy as np
+import sys
 from dotenv import load_dotenv
+import random
 from huggingface_hub import InferenceClient
 
 # Load environment variables
@@ -12,12 +15,13 @@ load_dotenv()
 MAX_TOKENS = 4000
 DEFAULT_TEMPERATURE = 0.5
 
-# Initialize the client
-client = OpenAI(
-    base_url="https://api-inference.huggingface.co/v1",
-    api_key=os.environ.get('API_KEY')  # Replace with your token
-)
+# initialize the client
 
+client = InferenceClient(
+  base_url="https://api-inference.huggingface.co/v1",
+  api_key=os.environ.get('API_KEY')  # Replace with your token
+)
+    
 # Create supported models
 model_links = {
     "Meta-Llama-3.1-8B": "meta-llama/Meta-Llama-3.1-8B-Instruct",
@@ -26,6 +30,7 @@ model_links = {
     "Falcon-7b-Instruct": "tiiuae/falcon-7b-instruct",
 }
 
+
 # Random dog images for error message
 random_dog_images = ["broken_llama3.jpeg"]
 
@@ -37,11 +42,12 @@ def reset_conversation():
     st.session_state.messages = []
     return None
 
-st.sidebar.button('Reset Chat', on_click=reset_conversation)  # Reset button
-
+st.sidebar.button('Reset Chat', on_click=reset_conversation) #Reset button
+    
 def main():
     st.header('Multi-Models')
 
+
     # Sidebar for model selection and temperature
     selected_model = st.sidebar.selectbox("Select Model", list(model_links.keys()))
     temperature = st.sidebar.slider('Select a temperature value', 0.0, 1.0, DEFAULT_TEMPERATURE)
@@ -51,6 +57,7 @@ def main():
 
     if st.session_state.prev_option != selected_model:
         st.session_state.messages = []
+        # st.write(f"Changed to {selected_model}")
         st.session_state.prev_option = selected_model
         reset_conversation()
 
@@ -60,11 +67,14 @@ def main():
     st.sidebar.write(f"You're now chatting with **{selected_model}**")
     st.sidebar.markdown("*Generated content may be inaccurate or false.*")
 
-    # Add a placeholder for chat messages
-    if "messages" in st.session_state:
-        for message in st.session_state.messages:
-            with st.chat_message(message["role"]):
-                st.markdown(message["content"])
+    # 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"])
 
     # Chat input and response
     if prompt := st.chat_input("Type message here..."):
@@ -78,14 +88,32 @@ def process_user_input(client, prompt, selected_model, temperature):
 
     # Generate and display assistant response
     with st.chat_message("assistant"):
-        response = """😵‍💫 Looks like someone unplugged something!
-        \n Either the model space is being updated or something is down."""
-        st.write(response)
-        random_dog_pick = random.choice(random_dog_images)
-        st.image(random_dog_pick)
-        st.write("This was the error message:")
-        # Replace with the actual error handling
-        st.write("Error details would go here.")
+        try:
+            stream = client.chat.completions.create(
+                model=model_links[selected_model],
+                messages=[
+                    {"role": m["role"], "content": m["content"]}
+                    for m in st.session_state.messages
+                ],
+                temperature=temperature,
+                stream=True,
+                max_tokens=MAX_TOKENS,
+            )
+            response = st.write_stream(stream)
+        except Exception as e:
+            handle_error(e)
+            return
+
+    st.session_state.messages.append({"role": "assistant", "content": response})
+
+def handle_error(error):
+    response = """😵‍💫 Looks like someone unplugged something!
+    \n Either the model space is being updated or something is down."""
+    st.write(response)
+    random_dog_pick = random.choice(random_dog_images)
+    st.image(random_dog_pick)
+    st.write("This was the error message:")
+    st.write(str(error))
 
 if __name__ == "__main__":
     main()