add wandb logging for user feedback

app.py

@@ -1,33 +1,24 @@
+import os
+import wandb
 import streamlit as st
+from constants import PROJECT_NAME
 from transformers import LayoutLMv3Processor, LayoutLMv3ForSequenceClassification
 from pdf2image import convert_from_bytes
 from PIL import Image
 
-labels = [
-    'budget',
-    'email',
-    'form',
-    'handwritten',
-    'invoice',
-    'language',
-    'letter',
-    'memo',
-    'news article',
-    'questionnaire',
-    'resume',
-    'scientific publication',
-    'specification',
-]
-id2label = {i: label for i, label in enumerate(labels)}
-label2id = {v: k for k, v in id2label.items()}
+wandb_api_key = os.getnev("WANDB_API_KEY")
+if not wandb_api_key:
+    st.error(
+        "Couldn't find WanDB API key. Please set it up as an environemnt variable",
+        icon="🚨",
+    )
+else:
+    wandb.login(key=wandb_api_key)
 
-processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-base")
-model = LayoutLMv3ForSequenceClassification.from_pretrained(
-    "microsoft/layoutlmv3-base",
-    num_labels=len(labels),
-    id2label=id2label,
-    label2id=label2id,
-)
+processor = LayoutLMv3Processor.from_pretrained("model/layoutlmv3/")
+model = LayoutLMv3ForSequenceClassification.from_pretrained("model/layoutlmv3/")
+id2label = model.config.id2label
+label2id = model.config.label2id
 
 st.title("Document Classification with LayoutLMv3")
 
@@ -36,6 +27,8 @@ uploaded_file = st.file_uploader(
 )
 
 if uploaded_file:
+    run = wandb.init(project=PROJECT_NAME, name='feedback-loop')
+
     if uploaded_file.type == "application/pdf":
         images = convert_from_bytes(uploaded_file.getvalue())
     else:
@@ -61,9 +54,21 @@ if uploaded_file:
             "Is the classification correct?", ("Yes", "No"),
             key=f'prediction-{i}'
         )
+
         if feedback == "No":
             correct_label = st.selectbox(
                 "Please select the correct label:", labels,
                 key=f'selectbox-{i}'
             )
             print(f'Correct label for image {i}: {correct_label}')
+
+            run.log({
+                'filepath': uploaded_file,
+                'filetype': uploaded_file.type,
+                'predicted_label': id2label[prediction],
+                'predicted_label_id': prediction,
+                'correct_label': correct_label,
+                'correct_label_id': label2id[correct_label]
+            })
+
+    run.finish()