Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import torch
 import streamlit as st
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import login
 
 # Access Hugging Face API token from environment variables
@@ -28,10 +28,8 @@ try:
     )
 
     # Define a chat template if it's missing
-    if not tokenizer.chat_template:
+    if not hasattr(tokenizer, "chat_template"):
         tokenizer.chat_template = """<s>[INST] {prompt} [/INST]"""
-
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 except Exception as e:
     st.error(f"Error loading model: {e}")
     st.stop()
@@ -40,17 +38,54 @@ except Exception as e:
 st.title("Pirate Chatbot")
 st.write("Ask me anything, and I'll respond in pirate speak!")
 
-# User input
-user_input = st.text_input("Your question:", "")
+# 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.')
+
+# Store conversation messages
+if "messages" not in st.session_state:
+    st.session_state.messages = [{"role": "assistant", "content": "Ahoy, matey! How can I assist ye?"}]
+
+# Display conversation
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.write(message["content"])
 
-if user_input:
-    # Format input using chat template
+def clear_chat_history():
+    st.session_state.messages = [{"role": "assistant", "content": "Ahoy, matey! How can I assist ye?"}]
+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}")
-    st.write(f"Formatted Prompt: {prompt}")  # Debugging: Print the formatted prompt
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     
     try:
-        outputs = pipe(prompt, max_new_tokens=256, do_sample=True)
-        st.write(outputs[0]["generated_text"])
+        # Generate response from the model
+        with torch.no_grad():
+            outputs = model.generate(inputs["input_ids"], max_new_tokens=256, do_sample=True)
+        
+        # Decode the generated response
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response
     except Exception as e:
         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})
+    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)
 
+        # Store the assistant's response
+        message = {"role": "assistant", "content": response}
+        st.session_state.messages.append(message)