Upload clean-chat-interface.py

clean-chat-interface.py

@@ -0,0 +1,103 @@
+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 Methods Engineer Assistant", page_icon="🔍", layout="wide")
+        self.initialize_session_state()
+        self.model_handler = self.load_model()
+
+    def initialize_session_state(self):
+        if "messages" not in st.session_state:
+            st.session_state.messages = [
+                {"role": "system", "content": "You are an experienced inspection methods engineer. Your task is to classify the following scope: analyze the scope provided in the input and determine the class item as an output."}
+            ]
+
+    @staticmethod
+    @st.cache_resource
+    def load_model():
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        st.info(f"Using device: {device}")
+        model_name = "amiguel/classItem-FT-llama-3-1-8b-instruct"
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            device_map="auto",
+            load_in_8bit=device == "cuda"
+        )
+        return ModelHandler(model, tokenizer)
+
+    def display_message(self, role, content):
+        with st.chat_message(role):
+            st.markdown(content)
+
+    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)
+        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")
+        
+        for message in st.session_state.messages:
+            if message["role"] != "system":
+                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})
+
+            conversation = "\n\n".join([msg["content"] for msg in st.session_state.messages])
+
+            with st.spinner("Analyzing and classifying scope..."):
+                response = self.model_handler.generate_response(conversation.strip())
+
+            clean_response = self.clean_response(response)
+            with st.chat_message("assistant"):
+                full_response = self.stream_response(clean_response)
+            st.session_state.messages.append({"role": "assistant", "content": full_response})
+
+        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}")
+
+    def clean_response(self, response):
+        # Remove any system: or user: prefixes from the response
+        lines = response.split('\n')
+        clean_lines = [line.split(':', 1)[-1].strip() if ':' in line else line for line in lines]
+        return '\n'.join(clean_lines)
+
+class ModelHandler:
+    def __init__(self, model, tokenizer):
+        self.model = model
+        self.tokenizer = tokenizer
+
+    def generate_response(self, conversation):
+        inputs = self.tokenizer(conversation, return_tensors="pt").to(self.model.device)
+        outputs = self.model.generate(**inputs, max_new_tokens=100)
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+if __name__ == "__main__":
+    app = ChatApp()
+    app.run()