Rename news_category_prediction.py to news_category_similar_news_prediction.py

news_category_prediction.py → news_category_similar_news_prediction.py

@@ -6,6 +6,7 @@ from config import (DISTILBERT_TOKENIZER_N_TOKENS,
                     CLASSIFIER_THRESHOLD)
 
 from logger import get_logger
+from find_similar_news import find_similar_news
 
 logger = get_logger()
 
@@ -43,7 +44,8 @@ def cols_check(new_cols, old_cols):
     return all([new_col==old_col for new_col, old_col in zip(new_cols, old_cols)])
 
     
-def predict_news_category(old_news: pd.DataFrame, new_news: pd.DataFrame, interpreter, label_encoder, tokenizer):
+def predict_news_category_similar_news(old_news: pd.DataFrame, new_news: pd.DataFrame, interpreter, label_encoder, tokenizer,
+                         collection, vectorizer, sent_model, ce_model):
     try:
         db_updation_required = 1
         logger.warning('Entering predict_news_category()')
@@ -73,6 +75,8 @@ def predict_news_category(old_news: pd.DataFrame, new_news: pd.DataFrame, interp
             final_df.drop_duplicates(subset='url', keep='first', inplace=True)
             headlines = [*final_df['title']].copy()
             label, prob = inference(headlines, interpreter, label_encoder, tokenizer)
+            sent_embs = vectorizer.vectorize(headlines)
+            sim_news = [find_similar_news(text, collection, vectorizer, sent_model, ce_model) for text in sent_embs]
             final_df['category'] = label
             final_df['pred_proba'] = prob
             final_df.reset_index(drop=True, inplace=True)