Update news_category_similar_news_prediction.py

news_category_similar_news_prediction.py

@@ -76,22 +76,26 @@ def predict_news_category_similar_news(old_news: pd.DataFrame, new_news: pd.Data
             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]
+            sim_news = [find_similar_news(search_vec, collection, vectorizer, sent_model, ce_model) for search_vec in sent_embs]
             final_df['category'] = label
             final_df['pred_proba'] = prob
+            final_df['similar_news'] = sim_news
             final_df.reset_index(drop=True, inplace=True)
             final_df.loc[final_df['pred_proba']<CLASSIFIER_THRESHOLD, 'category'] = 'OTHERS'
         else:
             logger.warning('Prior predictions found in old news')
-            if not cols_check([*new_news.columns], [*old_news.columns][:-2]):
+            if not cols_check([*new_news.columns], [*old_news.columns][:-3]):
                 raise Exception("New and old cols don't match")
             old_urls = [*old_news['url']]
             new_news = new_news.loc[new_news['url'].isin(old_urls) == False, :]
             if len(new_news) > 0:
                 headlines = [*new_news['title']].copy()
                 label, prob = inference(headlines, interpreter, label_encoder, tokenizer)
+                sent_embs = vectorizer.vectorize(headlines)
+                sim_news = [find_similar_news(search_vec, collection, vectorizer, sent_model, ce_model) for search_vec in sent_embs]
                 new_news['category'] = label
                 new_news['pred_proba'] = prob
+                final_df['similar_news'] = sim_news
                 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)