Update app.py

app.py

@@ -1,10 +1,13 @@
 import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import time
+import json
+from datetime import datetime
 
 class ChatApp:
     def __init__(self):
-        st.set_page_config(page_title="Inspection Engineer Chat", page_icon="🔍")
+        st.set_page_config(page_title="Inspection Engineer Chat", page_icon="🔍", layout="wide")
         self.initialize_session_state()
         self.model_handler = self.load_model()
 
@@ -36,24 +39,52 @@ class ChatApp:
     def get_user_input(self):
         return st.chat_input("Type your message here...")
 
+    def stream_response(self, response):
+        placeholder = st.empty()
+        full_response = ""
+        for word in response.split():
+            full_response += word + " "
+            placeholder.markdown(full_response + "▌")
+            time.sleep(0.01)  # Adjust the sleep time between 0.01 and 0.05 for desired speed
+        placeholder.markdown(full_response)
+        return full_response
+
+    def save_chat_history(self):
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        filename = f"chat_history_{timestamp}.json"
+        with open(filename, "w") as f:
+            json.dump(st.session_state.messages, f, indent=2)
+        return filename
+
     def run(self):
-        st.title("Inspection Methods Engineer Assistant")
+        col1, col2 = st.columns([3, 1])
+        
+        with col1:
+            st.title("Inspection Methods Engineer Assistant")
+
+            for message in st.session_state.messages[1:]:
+                self.display_message(message["role"], message["content"])
 
-        for message in st.session_state.messages[1:]:
-            self.display_message(message["role"], message["content"])
+            user_input = self.get_user_input()
+            if user_input:
+                self.display_message("user", user_input)
+                st.session_state.messages.append({"role": "user", "content": user_input})
 
-        user_input = self.get_user_input()
-        if user_input:
-            self.display_message("user", user_input)
-            st.session_state.messages.append({"role": "user", "content": user_input})
+                conversation = "\n".join([f"{msg['role']}: {msg['content']}" for msg in st.session_state.messages])
+                
+                with st.spinner("Analyzing and classifying scope..."):
+                    response = self.model_handler.generate_response(conversation)
+                
+                with st.chat_message("assistant"):
+                    full_response = self.stream_response(response)
+                
+                st.session_state.messages.append({"role": "assistant", "content": full_response})
 
-            conversation = "\n".join([f"{msg['role']}: {msg['content']}" for msg in st.session_state.messages])
-            
-            with st.spinner("Analyzing and classifying scope..."):
-                response = self.model_handler.generate_response(conversation)
-            
-            self.display_message("assistant", response)
-            st.session_state.messages.append({"role": "assistant", "content": response})
+        with col2:
+            st.sidebar.title("Chat Options")
+            if st.sidebar.button("Save Chat History"):
+                filename = self.save_chat_history()
+                st.sidebar.success(f"Chat history saved to {filename}")
 
 class ModelHandler:
     def __init__(self, model, tokenizer):