Update app.py

app.py

@@ -1,18 +1,17 @@
-import gradio as gr
-import pandas as pd
+ import gradio as gr
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
 import pickle
 
-# Existing code for loading model and tokenizer
+# Load the model and tokenizer from Hugging Face Hub
 model_name = "peterkros/cofogv1-bert"
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
+# Load the label encoder
 with open('label_encoder.pkl', 'rb') as file:
     label_encoder = pickle.load(file)
 
-# Existing prediction function
 def predict(text):
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
     with torch.no_grad():
@@ -22,35 +21,33 @@ def predict(text):
     predicted_label = label_encoder.inverse_transform([predicted_class])[0]
     return predicted_label
 
-# Function to display the metrics table
-def show_table():
-    data = {
-        "Epoch": [1, 2, 3, 4, 5],
-        "Training Loss": ["No log", "No log", "No log", "No log", "No log"],
-        "Validation Loss": [2.095209, 1.419945, 0.683810, 0.460408, 0.422096],
-        "Accuracy": [0.340764, 0.662420, 0.850318, 0.872611, 0.888535]
-    }
-    df = pd.DataFrame(data)
-    return df
-
-# Existing markdown text
+# Define the markdown text with bullet points
 markdown_text = """
 - Trained with ~1500 rows of data on bert-large-uncased, English.
 - Input one budget line per time.
 - Accuracy of the model is ~88%.
 """
+markdown_table = """
+| Epoch | Training Loss | Validation Loss | Accuracy  |
+|-------|---------------|-----------------|-----------|
+| 1     | No log        | 2.095209        | 0.340764  |
+| 2     | No log        | 1.419945        | 0.662420  |
+| 3     | No log        | 0.683810        | 0.850318  |
+| 4     | No log        | 0.460408        | 0.872611  |
+| 5     | No log        | 0.422096        | 0.888535  |
+"""
 
-# Update the Gradio Interface to include the table
 iface = gr.Interface(
-    fn=[predict, show_table], 
-    inputs=[gr.components.Textbox(lines=1, placeholder="Enter Budget line here...", label="Budget Input"), None], 
-    outputs=[gr.components.Textbox(label="Classification Output"), gr.outputs.DataFrame()],
+    fn=predict, 
+    inputs=gr.components.Textbox(lines=1, placeholder="Enter Budget line here...", label="Budget Input"), 
+    outputs=[gr.components.Textbox(label="Classification Output"), gr.components.Article(markdown_table)],
     title="COFOG Level 1 Classification",
     description=markdown_text,
-    allow_flagging="auto"
+    allow_flagging="auto"  # Enables flagging
 )
 
 # Run the interface
 if __name__ == "__main__":
     iface.launch()
 
+