Enhance SemanticSearcher integration and refine UpvotePredictor output handling

app.py

@@ -5,6 +5,7 @@ import pandas as pd
 from datasets import load_dataset
 
 from src.genai import GenAI
+from src.semantic_searcher import SemanticSearcher
 from src.upvote_predictor import UpvotePredictor
 
 # Load the dataset
@@ -52,18 +53,21 @@ examples = "\n".join(
 # Initialize the SemanticSearcher
 genai = GenAI()
 upvote_predictor = UpvotePredictor("src/bert_model")
-# _ = SemanticSearcher(df_counsel_chat_topic)
+ss = SemanticSearcher(df_counsel_chat_topic, df_counsel_chat)
 
 
 def get_output(question: str, question_context: str = None) -> str:
     answer, topic = genai.generate_content(
         question, question_context, unique_topics, examples
     )
-    # return (answer, topic, "Yes", pd.DataFrame())
     upvote_prediction = upvote_predictor.get_upvote_prediction(
         question, answer, question_context
     )
-    return (answer, topic, upvote_prediction[0], upvote_prediction[1])
+    if "not" in upvote_prediction.lower():
+        df = ss.retrieve_relevant_qna(question, question_context)
+        return (answer, topic, upvote_prediction, df)
+    else:
+        return (answer, topic, upvote_prediction, pd.DataFrame())
 
 
 demo = gr.Interface(

src/semantic_searcher.py

@@ -3,14 +3,16 @@ from sentence_transformers import SentenceTransformer
 
 
 class SemanticSearcher:
-    def __init__(self, df_counsel_chat_topic):
+    def __init__(self, df_counsel_chat_topic, df_counsel_chat):
         self.df_counsel_chat_topic = df_counsel_chat_topic
+        self.df_counsel_chat = df_counsel_chat
         self.embedder = SentenceTransformer("all-MiniLM-L6-v2")
         self.question_embeddings = self.embedder.encode(
             self.df_counsel_chat_topic["questionCombined"].tolist(),
             show_progress_bar=True,
             convert_to_tensor=True,
         )
+
     def retrieve_relevant_qna(self, question: str, question_context: str = None):
         if question_context is None:
             question_context = ""

src/upvote_predictor.py

@@ -1,5 +1,4 @@
 import numpy as np
-import pandas as pd
 import torch
 from torch.utils.data import DataLoader, SequentialSampler, TensorDataset
 from transformers import BertTokenizer
@@ -61,9 +60,6 @@ class UpvotePredictor:
                 predictions.extend(list(pred_flat))
 
         if predictions[0] == 0:
-            return (
-                "Not credible suggestion",
-                pd.DataFrame(),
-            )
+            return "Not credible suggestion"
         else:
-            return ("Credible suggestion", pd.DataFrame())
+            return "Credible suggestion"