Update gradio_app.py

gradio_app.py

@@ -112,6 +112,23 @@ def preprocess_and_flatten(json_results, mode, meta_fields=None):
     # Capitalize column names
     dataframe_results.columns = [col.capitalize() for col in dataframe_results.columns]
 
+    # Rename columns at the end of the function, conditionally if they exist
+    rename_columns = {}
+    
+    # Conditionally add renaming based on the mode and column existence
+    if 'doc_id' in dataframe_results.columns:
+        rename_columns['doc_id'] = 'DOC ID'
+    
+    if mode == 'claim_verification':
+        if 'doi' in dataframe_results.columns:
+            rename_columns['doi'] = 'DOI'
+        if 'sentence' in dataframe_results.columns:
+            rename_columns['sentence'] = 'Sentence'
+    
+    # Apply the renaming if there are any columns to rename
+    if rename_columns:
+        dataframe_results.rename(columns=rename_columns, inplace=True)
+
     return dataframe_results
     
 # Define the functions to handle the inputs and outputs