Update app.py

app.py

@@ -4,6 +4,7 @@ import pandas as pd
 import torch
 import os
 import re
+import base64  # For CSV download
 
 try:
     from transformers import AutoTokenizer, AutoModelForSequenceClassification
@@ -16,6 +17,19 @@ except ImportError as e:
 # Set page configuration
 st.set_page_config(page_title="WizNerd Insp", page_icon="🚀", layout="centered")
 
+# Custom CSS for Tw Cen MT font
+st.markdown("""
+    <style>
+    @import url('https://fonts.googleapis.com/css2?family=Tw+Cen+MT&display=swap');
+    html, body, [class*="css"]  {
+        font-family: 'Tw Cen MT', sans-serif !important;
+    }
+    .stTable table {
+        font-family: 'Tw Cen MT', sans-serif !important;
+    }
+    </style>
+""", unsafe_allow_html=True)
+
 # Load Hugging Face token
 HF_TOKEN = os.getenv("HF_TOKEN")
 
@@ -131,6 +145,13 @@ def classify_instruction(prompt, context, model, tokenizer):
             prediction = outputs.logits.argmax().item()
         return LABEL_TO_CLASS[prediction]
 
+# CSV download function
+def get_csv_download_link(df, filename="predicted_classes.csv"):
+    csv = df.to_csv(index=False)
+    b64 = base64.b64encode(csv.encode()).decode()  # Encode to base64
+    href = f'<a href="data:file/csv;base64,{b64}" download="{filename}">Download CSV</a>'
+    return href
+
 # Load model
 if "model" not in st.session_state:
     model_data = load_model(HF_TOKEN)
@@ -179,10 +200,11 @@ if prompt := st.chat_input("Ask your inspection question..."):
                     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 = file_data["content"][["Scope", "Functional Location"]].copy()
                         result_df["Predicted Class"] = predictions
                         st.write("Predicted Item Classes:")
                         st.table(result_df)
+                        st.markdown(get_csv_download_link(result_df), unsafe_allow_html=True)
                         response = "Predictions completed for uploaded file."
                     else:
                         predicted_class = classify_instruction(prompt, file_data["content"], model, tokenizer)