Update app.py

app.py

@@ -1,18 +1,34 @@
+import os
 import time 
 import gradio as gr
 import pandas as pd
 
+from classifier import classify
 
-def load_and_analyze_csv(file, text_field):
-    df = pd.read_csv(file.name)
+
+
+HFTOKEN = os.environ["HF_TOKEN"]
+
+def load_and_analyze_csv(file, text_field, event_model):
+    df = pd.read_table(file.name)
     
     if text_field not in df.columns:
         raise gr.Error(f"Error: Enter text column'{text_field}' not in CSV file.")
-    
-    fire_related = gr.CheckboxGroup(choices=df['text'].to_list()[:5])
-    flood_related = gr.CheckboxGroup(choices=df['text'].to_list()[:7])
-    not_related = gr.CheckboxGroup(choices=df['text'].to_list())
-    time.sleep(5)
+
+    floods, fires, nones = [], [], []
+    for post in df[text_field].to_list():
+        res = classify(post, event_model, HFTOKEN)
+        if res["event"] == 'flood':
+            floods.append(post)
+        elif res["event"] == 'fire':
+            fires.append(post)
+        else:
+            nones.append(post)
+            
+    fire_related = gr.CheckboxGroup(choices=fires)
+    flood_related = gr.CheckboxGroup(choices=floods)
+    not_related = gr.CheckboxGroup(choices=nones)
+    # time.sleep(5)
     return fire_related, flood_related, not_related
 
 def analyze_selected_texts(selections):
@@ -33,7 +49,7 @@ with gr.Blocks() as demo:
                 file_input = gr.File(label="Upload CSV File")
                 
             with gr.Column(scale=6):
-                text_field = gr.Textbox(label="Text field name", value="text")
+                text_field = gr.Textbox(label="Text field name", value="tweet_text")
                 event_model = gr.Dropdown(event_models, label="Select classification model")
                 predict_button = gr.Button("Start Prediction")
         
@@ -50,7 +66,8 @@ with gr.Blocks() as demo:
                 gr.Markdown("""### None""")
                 none_checkbox_output = gr.CheckboxGroup(label="Select ONLY incorrect classifications") 
                 
-            predict_button.click(load_and_analyze_csv, inputs=[file_input, text_field], outputs=[fire_checkbox_output, flood_checkbox_output, none_checkbox_output])           
+            predict_button.click(load_and_analyze_csv, inputs=[file_input, text_field, event_model], 
+                                 outputs=[fire_checkbox_output, flood_checkbox_output, none_checkbox_output])           
                 
     with gr.Tab("Question Answering"):
         # XXX Add some button disabling here, if the classification process is not completed first XXX