Upload 3 files

chat-interface.py

@@ -0,0 +1,12 @@
+import streamlit as st
+
+class ChatInterface:
+    def __init__(self):
+        self.chat_input_key = "chat_input"
+
+    def get_user_input(self):
+        return st.chat_input("Type your message here...", key=self.chat_input_key)
+
+    def display_message(self, role, content):
+        with st.chat_message(role):
+            st.write(content)

model-handler.py

@@ -0,0 +1,38 @@
+import torch
+import time
+
+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", truncation=True, max_length=1024)
+        
+        start_time = time.time()
+        output = ""
+        
+        with torch.no_grad():
+            for _ in range(150):  # Increased range for potentially longer responses
+                generated = self.model.generate(
+                    **inputs,
+                    max_new_tokens=1,
+                    do_sample=True,
+                    top_k=50,
+                    top_p=0.95
+                )
+                
+                new_token = generated[0, -1].item()
+                new_word = self.tokenizer.decode([new_token])
+                output += new_word
+                
+                inputs = self.tokenizer(conversation + output, return_tensors="pt", truncation=True, max_length=1024)
+                
+                if time.time() - start_time >= 0.01:
+                    yield output
+                    start_time = time.time()
+                
+                if new_token == self.tokenizer.eos_token_id:
+                    break
+        
+        return output.strip()

streamlit-app.py

@@ -0,0 +1,55 @@
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+from chat_interface import ChatInterface
+from model_handler import ModelHandler
+
+# Set page configuration
+st.set_page_config(page_title="Inspection Engineer Chat", page_icon="🔍")
+
+# Initialize session state
+if "messages" not in st.session_state:
+    st.session_state.messages = [
+        {"role": "system", "content": "You are an experienced senior inspection engineer. Your task is to analyze the scope provided in the input and determine the class item as an output."}
+    ]
+
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained("amiguel/classItem-FT-llama-3-1-8b-instruct")
+    model = AutoModelForCausalLM.from_pretrained("amiguel/classItem-FT-llama-3-1-8b-instruct")
+    return ModelHandler(model, tokenizer)
+
+def main():
+    st.title("Inspection Engineer Assistant")
+
+    # Load model
+    model_handler = load_model()
+
+    # Initialize chat interface
+    chat_interface = ChatInterface()
+
+    # Display chat messages
+    for message in st.session_state.messages[1:]:  # Skip the system message
+        chat_interface.display_message(message["role"], message["content"])
+
+    # Chat input
+    user_input = chat_interface.get_user_input()
+
+    if user_input:
+        # Add user message to chat history
+        st.session_state.messages.append({"role": "user", "content": user_input})
+        chat_interface.display_message("user", user_input)
+
+        # Prepare the full conversation context
+        conversation = "\n".join([f"{msg['role']}: {msg['content']}" for msg in st.session_state.messages])
+
+        # Generate response
+        with st.spinner("Analyzing..."):
+            response = model_handler.generate_response(conversation)
+
+        # Add assistant message to chat history
+        st.session_state.messages.append({"role": "assistant", "content": response})
+        chat_interface.display_message("assistant", response)
+
+if __name__ == "__main__":
+    main()