Update app.py

app.py

@@ -45,18 +45,38 @@ def call_chatgpt(prompt: str) -> str:
     return ans
 
 
-def ai_judge(prompt: str) -> float:
-    """
-    Uses the ChatGPT function to identify whether the content can answer the question
+# def ai_judge(prompt: str) -> float:
+#     """
+#     Uses the ChatGPT function to identify whether the content can answer the question
 
-    Args:
-        prompt: A string that represents the prompt
+#     Args:
+#         prompt: A string that represents the prompt
 
-    Returns:
-        float: A score
-    """
+#     Returns:
+#         float: A score
+#     """
+
+#     return call_chatgpt(prompt)
 
-    return call_chatgpt(prompt)
+
+def ai_judge(sentence1: str, sentence2: str) -> float:
+
+    HF_TOKEN = os.environ["HF_TOKEN"]
+    API_URL = "https://api-inference.huggingface.co/models/sentence-transformers/msmarco-distilbert-base-tas-b"
+    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
+    
+    def helper(payload):
+        response = requests.post(API_URL, headers=headers, json=payload)
+        return response.json()
+    
+    data = helper(
+        {
+            "inputs": {
+                "source_sentence": sentence1,
+                "sentences": [sentence2]
+            }
+        }
+    return data[0]
 
 
 def query(payload: Dict[str, Any]) -> Dict[str, Any]:
@@ -251,18 +271,18 @@ if prompt := st.chat_input("Tell me about YSA"):
         for i in range(final_ref.shape[0]):
             this_quest = question
             this_content = final_ref["answers"][i]
-            prompt_for_ai_judge = f"""
-                The user asked a question: {question}
+            # prompt_for_ai_judge = f"""
+            #     The user asked a question: {question}
 
-                We have found this content: {this_content}
+            #     We have found this content: {this_content}
 
-                From 0 to 10, rate how well the content answer the user's question. 
+            #     From 0 to 10, rate how well the content answer the user's question. 
 
-                Only produce a number from 0 to 10 while 10 being the best at answer user's question.
+            #     Only produce a number from 0 to 10 while 10 being the best at answer user's question.
 
-                If the content is a list of questions or not related to the user's question or it says inference endpoint is down, then you should say 0, because it does not answer user's question.
-            """
-            this_score = ai_judge(prompt_for_ai_judge)
+            #     If the content is a list of questions or not related to the user's question or it says inference endpoint is down, then you should say 0, because it does not answer user's question.
+            # """
+            this_score = ai_judge(question, this_content)
             independent_ai_judge_score.append(this_score)
 
         final_ref["ai_judge"] = independent_ai_judge_score