Update app.py

app.py

@@ -5,8 +5,6 @@ import PyPDF2
 import pandas as pd
 import torch
 import os
-import time
-import re
 
 # Set page configuration
 st.set_page_config(
@@ -39,10 +37,9 @@ BOT_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/99
 # Sidebar configuration
 with st.sidebar:
     st.header("Upload Documents 📂")
-    uploaded_files = st.file_uploader(
-        "Choose PDF, XLSX, or CSV files",
-        type=["pdf", "xlsx", "csv"],
-        accept_multiple_files=True,  # Allow multiple file uploads
+    uploaded_file = st.file_uploader(
+        "Choose a PDF or XLSX file",
+        type=["pdf", "xlsx"],
         label_visibility="collapsed"
     )
 
@@ -52,38 +49,20 @@ if "messages" not in st.session_state:
 
 # File processing function
 @st.cache_data
-def process_files(uploaded_files):
-    if not uploaded_files:
-        return []
+def process_file(uploaded_file):
+    if uploaded_file is None:
+        return ""
     
-    scopes = []
     try:
-        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()])
-    
+        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()
     except Exception as e:
         st.error(f"📄 Error processing file: {str(e)}")
-        return []
-    
-    return scopes
+        return ""
 
 # Model loading function
 @st.cache_resource
@@ -99,7 +78,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),
+            num_labels=len(LABEL_TO_CLASS),  # Ensure correct number of labels
             token=hf_token
         )
         
@@ -113,7 +92,7 @@ def load_model(hf_token):
         st.error(f"🤖 Model loading failed: {str(e)}")
         return None
 
-# Classification function with streaming simulation
+# Classification function
 def classify_instruction(prompt, file_context, model, tokenizer):
     full_prompt = f"Context:\n{file_context}\n\nInstruction: {prompt}"
     
@@ -130,17 +109,6 @@ 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:
@@ -170,35 +138,16 @@ if prompt := st.chat_input("Ask your inspection question..."):
         st.markdown(prompt)
     st.session_state.messages.append({"role": "user", "content": prompt})
 
-    # Process file context (if any)
-    file_scopes = process_files(uploaded_files)
-    file_context = "\n".join(file_scopes) if file_scopes else ""
+    # Process file context
+    file_context = process_file(uploaded_file)
     
-    # Classify the user prompt
+    # Classify the instruction
     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}"
-                
-                # 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.markdown(response)
                 st.session_state.messages.append({"role": "assistant", "content": response})
                 
         except Exception as e: