Update app.py

app.py

@@ -1,20 +1,22 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from huggingface_hub import login
 import PyPDF2
 import pandas as pd
 import torch
 import os
 import re
 
+try:
+    from transformers import AutoTokenizer, AutoModelForSequenceClassification
+    from huggingface_hub import login
+    TRANSFORMERS_AVAILABLE = True
+except ImportError as e:
+    st.error(f"Failed to import transformers: {str(e)}. Please install it with `pip install transformers`.")
+    TRANSFORMERS_AVAILABLE = False
+
 # Set page configuration
-st.set_page_config(
-    page_title="WizNerd Insp",
-    page_icon="🚀",
-    layout="centered"
-)
+st.set_page_config(page_title="WizNerd Insp", page_icon="🚀", layout="centered")
 
-# Load Hugging Face token from environment variable
+# Load Hugging Face token
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 # Model name
@@ -28,14 +30,14 @@ LABEL_TO_CLASS = {
     11: "Pressure Vessel (VII)", 12: "Structure", 13: "Flame Arrestor"
 }
 
-# Title with rocket emojis
+# Title
 st.title("🚀 WizNerd Insp 🚀")
 
-# Configure Avatars
+# Avatars
 USER_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/9904d9a0d445ab0488cf7395cb863cce7621d897/USER_AVATAR.png"
 BOT_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/991f4c6e4e1dc7a8e24876ca5aae5228bcdb4dba/Ataliba_Avatar.jpg"
 
-# Sidebar configuration
+# Sidebar
 with st.sidebar:
     st.header("Upload Documents 📂")
     uploaded_file = st.file_uploader(
@@ -44,11 +46,15 @@ with st.sidebar:
         label_visibility="collapsed"
     )
 
-# Initialize chat history
+# Initialize session state
 if "messages" not in st.session_state:
     st.session_state.messages = []
+if "file_processed" not in st.session_state:
+    st.session_state.file_processed = False
+if "file_data" not in st.session_state:
+    st.session_state.file_data = None
 
-# File processing function with pre-processing
+# File processing function
 @st.cache_data
 def process_file(uploaded_file):
     if uploaded_file is None:
@@ -58,24 +64,24 @@ def process_file(uploaded_file):
         if uploaded_file.type == "application/pdf":
             pdf_reader = PyPDF2.PdfReader(uploaded_file)
             text = "\n".join([page.extract_text() for page in pdf_reader.pages])
-            # Basic pre-processing
             text = re.sub(r'\s+', ' ', text.lower().strip())
             return {"type": "text", "content": text}
         
-        elif uploaded_file.type == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet":
-            df = pd.read_excel(uploaded_file)
-        elif uploaded_file.type == "text/csv":
-            df = pd.read_csv(uploaded_file)
-        
-        # For tabular data (xlsx, csv), detect scope columns
-        if 'df' in locals():
-            scope_cols = [col for col in df.columns if "scope" in col.lower()]
-            if not scope_cols:
-                st.warning("No 'scope' column found in the file. Using all data as context.")
-                return {"type": "table", "content": df.to_markdown()}
-            # Pre-process scope data
-            scope_data = df[scope_cols].dropna().astype(str).apply(lambda x: re.sub(r'\s+', ' ', x.lower().strip()))
-            return {"type": "scope", "content": scope_data}
+        elif uploaded_file.type in ["application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "text/csv"]:
+            df = pd.read_excel(uploaded_file) if uploaded_file.type == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" else pd.read_csv(uploaded_file)
+            required_cols = ["Scope", "Functional Location"]
+            available_cols = [col for col in required_cols if col in df.columns]
+            
+            if not available_cols:
+                st.warning("No 'Scope' or 'Functional Location' columns found. Treating as plain text.")
+                return {"type": "text", "content": df.to_string()}
+            
+            # Pre-process and concatenate Scope and Functional Location
+            df = df.dropna(subset=available_cols)
+            df["input_text"] = df[available_cols].apply(
+                lambda row: " ".join([re.sub(r'\s+', ' ', str(val).lower().strip()) for val in row]), axis=1
+            )
+            return {"type": "table", "content": df[["input_text"] + available_cols]}
         
     except Exception as e:
         st.error(f"📄 Error processing file: {str(e)}")
@@ -84,36 +90,31 @@ def process_file(uploaded_file):
 # Model loading function
 @st.cache_resource
 def load_model(hf_token):
+    if not TRANSFORMERS_AVAILABLE:
+        return None
     try:
         if not hf_token:
-            st.error("🔐 Authentication required! Please set the HF_TOKEN environment variable.")
+            st.error("🔐 Please set the HF_TOKEN environment variable.")
             return None
-        
         login(token=hf_token)
         tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=hf_token)
-        model = AutoModelForSequenceClassification.from_pretrained(
-            MODEL_NAME,
-            num_labels=len(LABEL_TO_CLASS),
-            token=hf_token
-        )
+        model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=len(LABEL_TO_CLASS), token=hf_token)
         device = "cuda" if torch.cuda.is_available() else "cpu"
         model.to(device)
         return model, tokenizer
-        
     except Exception as e:
         st.error(f"🤖 Model loading failed: {str(e)}")
         return None
 
 # Classification function
-def classify_instruction(prompt, file_context, model, tokenizer):
+def classify_instruction(prompt, context, model, tokenizer):
     model.eval()
     device = model.device
     
-    if file_context["type"] == "scope":
-        # Batch prediction for multiple scope entries
+    if isinstance(context, pd.DataFrame):
         predictions = []
-        for scope in file_context["content"].values.flatten():
-            full_prompt = f"Context:\n{scope}\n\nInstruction: {prompt}"
+        for text in context["input_text"]:
+            full_prompt = f"Context:\n{text}\n\nInstruction: {prompt}"
             inputs = tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True, max_length=128)
             inputs = {k: v.to(device) for k, v in inputs.items()}
             with torch.no_grad():
@@ -122,8 +123,7 @@ def classify_instruction(prompt, file_context, model, tokenizer):
                 predictions.append(LABEL_TO_CLASS[prediction])
         return predictions
     else:
-        # Single prediction for text or table context
-        full_prompt = f"Context:\n{file_context['content']}\n\nInstruction: {prompt}"
+        full_prompt = f"Context:\n{context}\n\nInstruction: {prompt}"
         inputs = tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True, max_length=128)
         inputs = {k: v.to(device) for k, v in inputs.items()}
         with torch.no_grad():
@@ -131,6 +131,29 @@ def classify_instruction(prompt, file_context, model, tokenizer):
             prediction = outputs.logits.argmax().item()
         return LABEL_TO_CLASS[prediction]
 
+# Load model
+if "model" not in st.session_state:
+    model_data = load_model(HF_TOKEN)
+    if model_data is None and TRANSFORMERS_AVAILABLE:
+        st.error("Failed to load model. Check HF_TOKEN.")
+        st.stop()
+    elif TRANSFORMERS_AVAILABLE:
+        st.session_state.model, st.session_state.tokenizer = model_data
+
+model = st.session_state.get("model")
+tokenizer = st.session_state.get("tokenizer")
+
+# Process uploaded file once
+if uploaded_file and not st.session_state.file_processed:
+    file_data = process_file(uploaded_file)
+    if file_data:
+        st.session_state.file_data = file_data
+        st.session_state.file_processed = True
+        if file_data["type"] == "table":
+            st.write("File uploaded with Scope and Functional Location data. Please provide an instruction.")
+        else:
+            st.write("File uploaded as text context. Please provide an instruction.")
+
 # Display chat messages
 for message in st.session_state.messages:
     avatar = USER_AVATAR if message["role"] == "user" else BOT_AVATAR
@@ -139,49 +162,38 @@ for message in st.session_state.messages:
 
 # Chat input handling
 if prompt := st.chat_input("Ask your inspection question..."):
-    # Load model if not already loaded
-    if "model" not in st.session_state:
-        model_data = load_model(HF_TOKEN)
-        if model_data is None:
-            st.error("Failed to load model. Please ensure HF_TOKEN is set correctly.")
-            st.stop()
-        st.session_state.model, st.session_state.tokenizer = model_data
-    
-    model = st.session_state.model
-    tokenizer = st.session_state.tokenizer
-    
+    if not TRANSFORMERS_AVAILABLE:
+        st.error("Transformers library not available.")
+        st.stop()
+
     # Add user message
     with st.chat_message("user", avatar=USER_AVATAR):
         st.markdown(prompt)
     st.session_state.messages.append({"role": "user", "content": prompt})
 
-    # Process file context
-    file_context = process_file(uploaded_file)
-    if file_context is None:
-        st.error("No file uploaded or file processing failed.")
-        st.stop()
-    
-    # Classify the instruction
+    # Handle response
     if model and tokenizer:
         try:
             with st.chat_message("assistant", avatar=BOT_AVATAR):
-                predicted_output = classify_instruction(prompt, file_context, model, tokenizer)
-                if file_context["type"] == "scope":
-                    # Display multiple predictions in a table
-                    scope_values = file_context["content"].values.flatten()
-                    result_df = pd.DataFrame({
-                        "Scope": scope_values,
-                        "Predicted Class": predicted_output
-                    })
-                    st.write("Predicted Classes:")
-                    st.table(result_df)
-                    response = "Predictions completed for multiple scope entries."
+                if st.session_state.file_data:
+                    file_data = st.session_state.file_data
+                    if file_data["type"] == "table":
+                        predictions = classify_instruction(prompt, file_data["content"], model, tokenizer)
+                        result_df = file_data["content"].copy()
+                        result_df["Predicted Class"] = predictions
+                        st.write("Predicted Item Classes:")
+                        st.table(result_df)
+                        response = "Predictions completed for uploaded file."
+                    else:
+                        predicted_class = classify_instruction(prompt, file_data["content"], model, tokenizer)
+                        response = f"The Item Class is: {predicted_class}"
                 else:
-                    # Single prediction
-                    response = f"The Item Class is: {predicted_output}"
-                    st.markdown(response)
-                st.session_state.messages.append({"role": "assistant", "content": response})
+                    # Handle single prompt without file
+                    predicted_class = classify_instruction(prompt, "", model, tokenizer)
+                    response = f"The Item Class is: {predicted_class}"
                 
+                st.markdown(response)
+                st.session_state.messages.append({"role": "assistant", "content": response})
         except Exception as e:
             st.error(f"⚡ Classification error: {str(e)}")
     else: