Update model_comparison.py

model_comparison.py

@@ -41,23 +41,21 @@ TASK_EVAL_DATABASE_PATH = 'data/task_oriented_eval_database.yaml'
 #         st.write("NEW EVAL ID:", newEvalID)
 #     return newEvalID
 
-# def check_profanity(df):
-#     cleanedDF = df
-#     for i, row in cleanedDF.iterrows():
-#         if predict([row['User']])[0] != 0.0:
-#             cleanedDF.at[i, 'User'] = '**NSFW**'
-#         if 'Target' in df:
-#             if predict([row['Target']])[0] != 0.0:
-#                 cleanedDF.at[i, 'Target'] = '**NSFW**'
-#     return cleanedDF
-
+def check_profanity(df):
+    cleanedDF = df
+    for i, row in cleanedDF.iterrows():
+        if 'Target' in df:
+            if predict([row['Target']])[0] != 0.0:
+                cleanedDF.at[i, 'Target'] = '**NSFW**'
+    return cleanedDF
 def dataframe_with_selections(df):
-    # df_with_selections = check_profanity(df.copy()) 
-    df.insert(0, "Select", True)
+    df_with_selections = df.copy()
+    df_with_selections = check_profanity(df_with_selections)
+    df_with_selections.insert(0, "Select", True)
 
     # Get dataframe row-selections from user with st.data_editor
     edited_df = st.data_editor(
-        df,
+        df_with_selections,
         hide_index=True,
         column_config={"Select": st.column_config.CheckboxColumn(required=True)},
         disabled=df.columns,