Update news_category_prediction.py

news_category_prediction.py

@@ -1,7 +1,13 @@
 import pandas as pd
 import numpy as np
 import tensorflow as tf 
-from config import DISTILBERT_TOKENIZER_N_TOKENS, NEWS_CATEGORY_CLASSIFIER_N_CLASSES, CLASSIFIER_THRESHOLD
+from config import (DISTILBERT_TOKENIZER_N_TOKENS, 
+                    NEWS_CATEGORY_CLASSIFIER_N_CLASSES, 
+                    CLASSIFIER_THRESHOLD)
+
+from logger import get_logger
+
+logger = get_logger()
 
 def parse_prediction(tflite_pred, label_encoder):
     tflite_pred_argmax = np.argmax(tflite_pred, axis=1)
@@ -11,6 +17,7 @@ def parse_prediction(tflite_pred, label_encoder):
     
 
 def inference(text, interpreter, label_encoder, tokenizer):
+    logger.warning('Entering inference()')
     batch_size = len(text)
     if text != "":
         tokens = tokenizer(text, max_length=DISTILBERT_TOKENIZER_N_TOKENS, padding="max_length", truncation=True, return_tensors="tf")
@@ -28,6 +35,7 @@ def inference(text, interpreter, label_encoder, tokenizer):
         interpreter.invoke()
         tflite_pred = interpreter.get_tensor(output_details["index"])
         tflite_pred = parse_prediction(tflite_pred)
+    logger.warning('Exiting inference()')
     return tflite_pred
 
 def cols_check(new_cols, old_cols):
@@ -35,6 +43,7 @@ def cols_check(new_cols, old_cols):
 
     
 def predict_news_category(old_news: pd.DataFrame, new_news: pd.DataFrame, interpreter, label_encoder, tokenizer):
+    logger.warning('Entering predicting_news_category()')
     old_news = old_news.copy()
     new_news = new_news.copy()
     # dbops = DBOperations()
@@ -42,7 +51,7 @@ def predict_news_category(old_news: pd.DataFrame, new_news: pd.DataFrame, interp
     old_news.drop(columns='_id', inplace=True)
     # new_news = get_news()  
     if 'category' not in [*old_news.columns]:
-        print('no prior predictions found')
+        logger.warning('No prior predictions found in old news')
         if not cols_check([*new_news.columns], [*old_news.columns]):
             raise Exeption("New and old cols don't match")
         final_df = pd.concat([old_news, new_news], axis=0, ignore_index=True)
@@ -52,7 +61,7 @@ def predict_news_category(old_news: pd.DataFrame, new_news: pd.DataFrame, interp
         final_df['category'] = label
         final_df['pred_proba'] = prob
     else:
-        print('prior predictions found')
+        logger.warning('Prior predictions found in old news')
         if not cols_check([*new_news.columns], [*old_news.columns][:-2]):
             raise Exeption("New and old cols don't match")
         old_urls = [*old_news['url']]
@@ -65,5 +74,6 @@ def predict_news_category(old_news: pd.DataFrame, new_news: pd.DataFrame, interp
         final_df.drop_duplicates(subset='url', keep='first', inplace=True)
     final_df.reset_index(drop=True, inplace=True)
     final_df.loc[final_df['pred_proba']<CLASSIFIER_THRESHOLD, 'category'] = 'OTHERS'
+    logger.warning('Exiting predicting_news_category()')
     return final_df