Generate a prediction table

app.py

@@ -1,8 +1,6 @@
 # TODO: requirments.txt
-# TODO: Secrets?
-# TODO: Dataset_name to Secrets and not hard-coded!
-
 import os
+import pandas as pd
 import streamlit as st
 
 import torch
@@ -12,7 +10,8 @@ from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 model_name = st.text_input("Enter a model's name on HF")
 # MODEL_NAME = "AMR-KELEG/NADI2024-baseline"
-DIALECTS = ["Algeria",
+DIALECTS = [
+    "Algeria",
     "Bahrain",
     "Egypt",
     "Iraq",
@@ -36,6 +35,7 @@ assert len(DIALECTS) == 18
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
+
 def predict_top_p(text, P=0.9):
     """Predict the top dialects with an accumulative confidence of at least P."""
     assert P <= 1 and P >= 0
@@ -56,12 +56,29 @@ def predict_top_p(text, P=0.9):
     return [DIALECTS[i] for i, p in enumerate(predictions) if p == 1]
 
 
+# Load the dataset
 dataset_name = "AMR-KELEG/test-dataset"
 dataset = datasets.load_dataset(dataset_name, token=os.environ["HF_TOKEN"])["test"]
 
+sentences_labels, sentences_predictions = [], []
+
 for sample in tqdm(dataset):
     text = sample["sentence"]
-    labels= [DIALECTS[i] for i in range(len(DIALECTS)) if DIALECTS[i] in sample.keys() and int(sample[DIALECTS[i]]) == 1]
+    labels = [
+        DIALECTS[i]
+        for i in range(len(DIALECTS))
+        if DIALECTS[i] in sample.keys() and int(sample[DIALECTS[i]]) == 1
+    ]
     pred = predict_top_p(text)
-    sample["pred"] = pred
-    st.write("Text:", text)
+    sentences_labels.append(labels)
+    sentences_predictions.append(pred)
+
+st.table(
+    data=pd.DataFrame(
+        {
+            "text": dataset["sentence"],
+            "labels": sentences_labels,
+            "predictions": sentences_predictions,
+        }
+    )
+)

requirements.txt

@@ -1,3 +1,4 @@
 transformers
 torch
 datasets
+pandas