Update app.py

app.py

@@ -1,70 +1,74 @@
+import streamlit as st
 import os
 import torch
-import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import login
 
-# Access Hugging Face API token from environment variables
-api_key = os.getenv("llama3")
-
-if not api_key:
-    st.error("Hugging Face API token is missing!")
-    st.stop()  # Stop execution if no API key is found
-
-# Authenticate with Hugging Face Hub
-try:
-    login(api_key)
-except Exception as e:
-    st.error(f"Authentication failed: {e}")
-    st.stop()
+# App title
+st.set_page_config(page_title="🦙💬 Llama 2 Chatbot")
 
-# Load the model and tokenizer
-model_id = "meta-llama/Llama-3.2-1B"
+# Hugging Face Credentials
+with st.sidebar:
+    st.title('🦙💬 Llama 2 Chatbot')
+    st.write('This chatbot is created using the open-source Llama model from Meta.')
+    
+    # Use Hugging Face API Key from secrets or environment
+    api_key = os.getenv("HUGGINGFACE_API_KEY")
 
-try:
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id, torch_dtype=torch.bfloat16, device_map="auto"
-    )
+    if not api_key:
+        st.error("Hugging Face API key is missing!")
+        st.stop()
 
-    # Define a chat template if it's missing
-    if not hasattr(tokenizer, "chat_template"):
-        tokenizer.chat_template = """<s>[INST] {prompt} [/INST]"""
-except Exception as e:
-    st.error(f"Error loading model: {e}")
-    st.stop()
+    # Authenticate with Hugging Face Hub
+    try:
+        login(api_key)
+        st.success('API key successfully authenticated!', icon='✅')
+    except Exception as e:
+        st.error(f"Authentication failed: {e}")
+        st.stop()
 
-# Streamlit interface
-st.title("Pirate Chatbot")
-st.write("Ask me anything, and I'll respond in pirate speak!")
+    st.subheader('Models and parameters')
+    selected_model = st.sidebar.selectbox('Choose a Llama model', ['Llama-3.2-1B', 'Llama-7B'], key='selected_model')
+    
+    # Model configurations based on selection
+    if selected_model == 'Llama-3.2-1B':
+        model_id = "meta-llama/Llama-3.2-1B"
+    elif selected_model == 'Llama-7B':
+        model_id = "meta-llama/Llama-7B"
 
-# Sidebar settings (removed API key input)
-with st.sidebar:
-    st.title('Pirate Chatbot')
-    st.write('This chatbot uses the Llama 2 model for chat. You can interact with it directly.')
+    temperature = st.sidebar.slider('temperature', min_value=0.01, max_value=1.0, value=0.1, step=0.01)
+    top_p = st.sidebar.slider('top_p', min_value=0.01, max_value=1.0, value=0.9, step=0.01)
+    max_length = st.sidebar.slider('max_length', min_value=20, max_value=80, value=50, step=5)
+    st.markdown('📖 Learn how to build this app in this [blog](https://blog.streamlit.io/how-to-build-a-llama-2-chatbot/)!')
 
-# Store conversation messages
-if "messages" not in st.session_state:
-    st.session_state.messages = [{"role": "assistant", "content": "Ahoy, matey! How can I assist ye?"}]
+# Store LLM generated responses
+if "messages" not in st.session_state.keys():
+    st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?"}]
 
-# Display conversation
+# Display or clear chat messages
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.write(message["content"])
 
 def clear_chat_history():
-    st.session_state.messages = [{"role": "assistant", "content": "Ahoy, matey! How can I assist ye?"}]
+    st.session_state.messages = [{"role": "assistant", "content": "How may I assist you today?"}]
 st.sidebar.button('Clear Chat History', on_click=clear_chat_history)
 
-# Generate response function
-def generate_pirate_response(user_input):
-    prompt = tokenizer.chat_template.format(prompt=f"You are a pirate chatbot! Answer in pirate speak!\nUser: {user_input}")
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    
+# Load the tokenizer and model
+try:
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
+except Exception as e:
+    st.error(f"Error loading model: {e}")
+    st.stop()
+
+# Function for generating response using Hugging Face model
+def generate_huggingface_response(prompt_input):
+    inputs = tokenizer(prompt_input, return_tensors="pt").to(model.device)
     try:
         # Generate response from the model
         with torch.no_grad():
-            outputs = model.generate(inputs["input_ids"], max_new_tokens=256, do_sample=True)
+            outputs = model.generate(inputs["input_ids"], max_new_tokens=max_length, temperature=temperature, top_p=top_p, do_sample=True)
         
         # Decode the generated response
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
@@ -73,19 +77,22 @@ def generate_pirate_response(user_input):
         st.error(f"Error generating response: {e}")
         return "Oops! Something went wrong."
 
-# Handle user input and response generation
-if user_input := st.text_input("Your question:"):
-    st.session_state.messages.append({"role": "user", "content": user_input})
+# User-provided prompt
+if prompt := st.chat_input(disabled=not api_key):
+    st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user"):
-        st.write(user_input)
-
-    # Generate and display assistant's response
-    if st.session_state.messages[-1]["role"] != "assistant":
-        with st.chat_message("assistant"):
-            with st.spinner("Thinking..."):
-                response = generate_pirate_response(user_input)
-                st.write(response)
+        st.write(prompt)
 
-        # Store the assistant's response
-        message = {"role": "assistant", "content": response}
-        st.session_state.messages.append(message)
+# Generate a new response if last message is not from assistant
+if st.session_state.messages[-1]["role"] != "assistant":
+    with st.chat_message("assistant"):
+        with st.spinner("Thinking..."):
+            response = generate_huggingface_response(prompt)
+            placeholder = st.empty()
+            full_response = ''
+            for item in response:
+                full_response += item
+                placeholder.markdown(full_response)
+            placeholder.markdown(full_response)
+    message = {"role": "assistant", "content": full_response}
+    st.session_state.messages.append(message)