Update news_category_similar_news_prediction.py

news_category_similar_news_prediction.py

@@ -105,8 +105,9 @@ def add_flags_to_prediction_df(df):
     df['text'] = headlines_desc
     df['used_for_training'] = 0
     df['allocated'] = 0
+    df.rename(columns={'category', 'y_pred', 'parsed_date': 'date_time'}, inplace=True)
     
-    df.drop(columns=['similar_news', 'parsed_date', 'src', 'title', 'description'], inplace=True)
+    df.drop(columns=['similar_news', 'src', 'title', 'description'], inplace=True)
     return df