Update news_category_prediction.py

news_category_prediction.py

@@ -43,46 +43,51 @@ 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 predict_news_category()')
-    old_news = old_news.copy()
-    new_news = new_news.copy()
-    logger.warning(f'old news columns: {[*old_news.columns]}')
-    logger.warning(f'new news columns: {[*new_news.columns]}')
-    if len(old_news) == 0:
-        logger.warning('No old news is found')
-        old_news = new_news.copy()
-    else:
-        logger.warning('Old news is found')
-        old_news.drop(columns='_id', inplace=True)
-
-    if len(new_news) <= 5:
-        raise Exception("No New News Found")
-
-    
-    if 'category' not in [*old_news.columns]:
-        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)
-        final_df.drop_duplicates(subset='url', keep='first', inplace=True)
-        headlines = [*final_df['title']].copy()
-        label, prob = inference(headlines, interpreter, label_encoder, tokenizer)
-        final_df['category'] = label
-        final_df['pred_proba'] = prob
-    else:
-        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']]
-        new_news = new_news.loc[new_news['url'].isin(old_urls) == False, :]
-        headlines = [*new_news['title']].copy()
-        label, prob = inference(headlines, interpreter, label_encoder, tokenizer)
-        new_news['category'] = label
-        new_news['pred_proba'] = prob
-        final_df = pd.concat([old_news, new_news], axis=0, ignore_index=True)
-        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 predict_news_category()')
+    try:
+        logger.warning('Entering predict_news_category()')
+        logger.warning(f'old news: {old_news}\nnew_news: {new_news}')
+        if not isinstance(new_news, pd.DataFrame):
+            raise Exception('No New News Found')
+        else:
+            new_news = new_news.copy()
+            logger.warning(f'new news columns: {[*new_news.columns]}')
+            logger.warning(f'{len(new_news)} new news items found')
+        
+        if isinstance(old_news, pd.DataFrame):
+            old_news = old_news.copy()
+            logger.warning(f'old news columns: {[*old_news.columns]}')
+            logger.warning(f'{len(old_news)} old news items found')
+        else:
+            logger.warning('No old news is found')
+            old_news = new_news.copy()
+      
+        if 'category' not in [*old_news.columns]:
+            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)
+            final_df.drop_duplicates(subset='url', keep='first', inplace=True)
+            headlines = [*final_df['title']].copy()
+            label, prob = inference(headlines, interpreter, label_encoder, tokenizer)
+            final_df['category'] = label
+            final_df['pred_proba'] = prob
+        else:
+            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']]
+            new_news = new_news.loc[new_news['url'].isin(old_urls) == False, :]
+            headlines = [*new_news['title']].copy()
+            label, prob = inference(headlines, interpreter, label_encoder, tokenizer)
+            new_news['category'] = label
+            new_news['pred_proba'] = prob
+            final_df = pd.concat([old_news, new_news], axis=0, ignore_index=True)
+            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 predict_news_category()')
+    except Exception as e:
+        logger.warning(f'Unexcpected error in predict_news_category()\n{e}')
+        return None
     return final_df