Update app.py

app.py

@@ -5,6 +5,8 @@ import PyPDF2
 import pandas as pd
 import torch
 import os
+import time
+import re
 
 # Set page configuration
 st.set_page_config(
@@ -37,9 +39,10 @@ BOT_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/99
 # Sidebar configuration
 with st.sidebar:
     st.header("Upload Documents 📂")
-    uploaded_file = st.file_uploader(
-        "Choose a PDF or XLSX file",
-        type=["pdf", "xlsx"],
+    uploaded_files = st.file_uploader(
+        "Choose PDF, XLSX, or CSV files",
+        type=["pdf", "xlsx", "csv"],
+        accept_multiple_files=True,  # Allow multiple file uploads
         label_visibility="collapsed"
     )
 
@@ -49,20 +52,38 @@ if "messages" not in st.session_state:
 
 # File processing function
 @st.cache_data
-def process_file(uploaded_file):
-    if uploaded_file is None:
-        return ""
+def process_files(uploaded_files):
+    if not uploaded_files:
+        return []
     
+    scopes = []
     try:
-        if uploaded_file.type == "application/pdf":
-            pdf_reader = PyPDF2.PdfReader(uploaded_file)
-            return "\n".join([page.extract_text() for page in pdf_reader.pages])
-        elif uploaded_file.type == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet":
-            df = pd.read_excel(uploaded_file)
-            return df.to_markdown()
+        for uploaded_file in uploaded_files:
+            if uploaded_file.type == "application/pdf":
+                pdf_reader = PyPDF2.PdfReader(uploaded_file)
+                text = "\n".join([page.extract_text() for page in pdf_reader.pages])
+                # Split text into potential scope lines (e.g., by newlines or sentences)
+                lines = [line.strip() for line in text.split("\n") if line.strip()]
+                # Filter lines that look like scope instructions (e.g., contain keywords like "at location", "DAL/")
+                scope_lines = [line for line in lines if re.search(r"(at location|DAL/|PSV-|CD-|DA-)", line, re.IGNORECASE)]
+                scopes.extend(scope_lines)
+            
+            elif uploaded_file.type in ["application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "text/csv"]:
+                if uploaded_file.type == "text/csv":
+                    df = pd.read_csv(uploaded_file)
+                else:
+                    df = pd.read_excel(uploaded_file)
+                # Assume the first column contains scope instructions
+                if not df.empty:
+                    scope_column = df.columns[0]  # First column
+                    scope_lines = df[scope_column].dropna().astype(str).tolist()
+                    scopes.extend([line.strip() for line in scope_lines if line.strip()])
+    
     except Exception as e:
         st.error(f"📄 Error processing file: {str(e)}")
-        return ""
+        return []
+    
+    return scopes
 
 # Model loading function
 @st.cache_resource
@@ -78,7 +99,7 @@ def load_model(hf_token):
         tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=hf_token)
         model = AutoModelForSequenceClassification.from_pretrained(
             MODEL_NAME,
-            num_labels=len(LABEL_TO_CLASS),  # Ensure correct number of labels
+            num_labels=len(LABEL_TO_CLASS),
             token=hf_token
         )
         
@@ -92,7 +113,7 @@ def load_model(hf_token):
         st.error(f"🤖 Model loading failed: {str(e)}")
         return None
 
-# Classification function
+# Classification function with streaming simulation
 def classify_instruction(prompt, file_context, model, tokenizer):
     full_prompt = f"Context:\n{file_context}\n\nInstruction: {prompt}"
     
@@ -109,6 +130,17 @@ def classify_instruction(prompt, file_context, model, tokenizer):
     
     return class_name
 
+def stream_classification_output(class_name, delay=0.05):
+    """Simulate streaming by displaying the class name character by character."""
+    response_container = st.empty()
+    full_response = ""
+    for char in class_name:
+        full_response += char
+        response_container.markdown(f"Predicted class: {full_response} ▌")
+        time.sleep(delay)
+    response_container.markdown(f"Predicted class: {full_response}")
+    return full_response
+
 # Display chat messages
 for message in st.session_state.messages:
     try:
@@ -138,16 +170,35 @@ if prompt := st.chat_input("Ask your inspection question..."):
         st.markdown(prompt)
     st.session_state.messages.append({"role": "user", "content": prompt})
 
-    # Process file context
-    file_context = process_file(uploaded_file)
+    # Process file context (if any)
+    file_scopes = process_files(uploaded_files)
+    file_context = "\n".join(file_scopes) if file_scopes else ""
     
-    # Classify the instruction
+    # Classify the user prompt
     if model and tokenizer:
         try:
             with st.chat_message("assistant", avatar=BOT_AVATAR):
+                # Classify the user-entered prompt
                 predicted_class = classify_instruction(prompt, file_context, model, tokenizer)
+                # Stream the classification output
+                streamed_response = stream_classification_output(predicted_class)
                 response = f"Predicted class: {predicted_class}"
-                st.markdown(response)
+                
+                # If there are scopes from files, classify them too
+                if file_scopes:
+                    st.markdown("### Classifications from Uploaded Files")
+                    results = []
+                    for scope in file_scopes:
+                        predicted_class = classify_instruction(scope, file_context, model, tokenizer)
+                        results.append({"Scope": scope, "Predicted Class": predicted_class})
+                    
+                    # Display results in a table
+                    df_results = pd.DataFrame(results)
+                    st.table(df_results)
+                    
+                    # Add table to chat history
+                    response += "\n\n### Classifications from Uploaded Files\n" + df_results.to_markdown(index=False)
+                
                 st.session_state.messages.append({"role": "assistant", "content": response})
                 
         except Exception as e: