Update retriever.py

retriever.py

@@ -1,121 +1,121 @@
-import pandas as pd
-import json
-import sys
-import os
-from collections import defaultdict
-from util.vector_base import EmbeddingFunction, get_or_create_vector_base
-from util.Embeddings import TextEmb3LargeEmbedding
-from langchain_core.documents import Document
-from FlagEmbedding import FlagReranker
-import time
-from bm25s import BM25, tokenize
-import contextlib
-import io
-from tqdm import tqdm
-
-def rrf(rankings, k = 60):
-    res = 0
-    for r in rankings:
-        res += 1 / (r + k)
-    return res
-
-def retriever(requirement, PO, safeguard_vector_store, reranker_model, using_reranker=False, using_BM25=False, using_chroma=True, k=20, if_split_po=True):
-    final_result = []
-    if not if_split_po:
-        final_result = multiretriever(requirement, PO, safeguard_vector_store, reranker_model, using_reranker=using_reranker, using_BM25=using_BM25, using_chroma=using_chroma, k=k)
-    else:
-        for po in PO:
-            po_result = multiretriever(requirement, [po], safeguard_vector_store, reranker_model, using_reranker=using_reranker, using_BM25=using_BM25, using_chroma=using_chroma, k=k)
-            for safeguard in po_result:
-                final_result.append(safeguard)
-    return final_result
-
-def multiretriever(requirement, PO, safeguard_vector_store, reranker_model, using_reranker=True, using_BM25=False, using_chroma=True, k=20):
-    """
-    requirements_dict: [
-        requirement: {
-            "PO": [],
-            "safeguard": []
-        }
-    ]
-    """
-    candidate_safeguards = []
-    po_list = [po.lower().rstrip() for po in PO if po]
-    if "young users" in po_list and len(po_list) == 1:
-        return []
-    candidate_safeguards = safeguard_vector_store.get(where={"po": {"$in": po_list}})
-    safeguard_dict, safeguard_content = {}, []
-    for id, content, metadata in zip(candidate_safeguards['ids'], candidate_safeguards['documents'], candidate_safeguards['metadatas']):
-        safeguard_dict[content] = {
-            "metadata": metadata,
-            "rank": [],
-            "rrf_score": 0
-        }
-        safeguard_content.append(content)
-     
-    # Reranker
-    if using_reranker:
-        content_pairs, reranking_rank, reranking_results = [], [], []
-        for safeguard in safeguard_content:
-            content_pairs.append([requirement, safeguard])
-        safeguard_rerank_scores = reranker_model.compute_score(content_pairs)
-        for content_pair, score in zip(content_pairs, safeguard_rerank_scores):
-            reranking_rank.append((content_pair[1], score))
-        reranking_results = sorted(reranking_rank, key=lambda x: x[1], reverse=True)
-        for safeguard, score in reranking_results:
-            safeguard_dict[safeguard]['rank'].append(reranking_results.index((safeguard, score)) + 1)
-
-    #  BM25
-    if using_BM25:
-        with contextlib.redirect_stdout(io.StringIO()):
-            bm25_retriever = BM25(corpus=safeguard_content)
-            bm25_retriever.index(tokenize(safeguard_content))
-            bm25_results, scores = bm25_retriever.retrieve(tokenize(requirement), k = len(safeguard_content))
-            bm25_retrieval_rank = 1
-            for safeguard in bm25_results[0]:
-                safeguard_dict[safeguard]['rank'].append(bm25_retrieval_rank)
-                bm25_retrieval_rank += 1
-
-        # chroma retrieval
-    if using_chroma:
-        retrieved_safeguards = safeguard_vector_store.similarity_search_with_score(query=requirement, k=len(candidate_safeguards['ids']), filter={"po": {"$in": po_list}})
-        retrieval_rank = 1
-        for safeguard in retrieved_safeguards:
-            safeguard_dict[safeguard[0].page_content]['rank'].append(retrieval_rank)
-            retrieval_rank += 1
-
-    final_result = []
-    for safeguard in safeguard_content:
-        safeguard_dict[safeguard]['rrf_score'] = rrf(safeguard_dict[safeguard]['rank'])
-        final_result.append((safeguard_dict[safeguard]['rrf_score'], safeguard_dict[safeguard]['metadata']['safeguard_number'], safeguard, safeguard_dict[safeguard]['metadata']['po']))
-    final_result.sort(key=lambda x: x[0], reverse=True)
-        
-        # top k 
-    topk_final_result = final_result[:k]
-
-    return topk_final_result
-
-if __name__=="__main__":
-    embeddingmodel = TextEmb3LargeEmbedding(max_qpm=58)
-    embedding = EmbeddingFunction(embeddingmodel)
-    safeguard_vector_store = get_or_create_vector_base('safeguard_database', embedding)
-    reranker_model = FlagReranker(
-        '/root/PTR-LLM/tasks/pcf/model/bge-reranker-v2-m3',
-        use_fp16=True,
-        devices=["cpu"],   
-    )
-    requirement = """
-    Data Minimization Consent for incompatible purposes: Require consent for additional use of personal information not reasonably necessary to or incompatible with original purpose disclosure.
-    """
-    PO = ["Data Minimization & Purpose Limitation", "Transparency"]
-    final_result = retriever(
-            requirement,
-            PO,
-            safeguard_vector_store,
-            reranker_model,
-            using_reranker=True,
-            using_BM25=False,
-            using_chroma=True,
-            k=10
-            )
+import pandas as pd
+import json
+import sys
+import os
+from collections import defaultdict
+from util.vector_base import EmbeddingFunction, get_or_create_vector_base
+from util.Embeddings import TextEmb3LargeEmbedding
+from langchain_core.documents import Document
+from FlagEmbedding import FlagReranker
+import time
+# from bm25s import BM25, tokenize
+import contextlib
+import io
+from tqdm import tqdm
+
+def rrf(rankings, k = 60):
+    res = 0
+    for r in rankings:
+        res += 1 / (r + k)
+    return res
+
+def retriever(requirement, PO, safeguard_vector_store, reranker_model, using_reranker=False, using_BM25=False, using_chroma=True, k=20, if_split_po=True):
+    final_result = []
+    if not if_split_po:
+        final_result = multiretriever(requirement, PO, safeguard_vector_store, reranker_model, using_reranker=using_reranker, using_BM25=using_BM25, using_chroma=using_chroma, k=k)
+    else:
+        for po in PO:
+            po_result = multiretriever(requirement, [po], safeguard_vector_store, reranker_model, using_reranker=using_reranker, using_BM25=using_BM25, using_chroma=using_chroma, k=k)
+            for safeguard in po_result:
+                final_result.append(safeguard)
+    return final_result
+
+def multiretriever(requirement, PO, safeguard_vector_store, reranker_model, using_reranker=True, using_BM25=False, using_chroma=True, k=20):
+    """
+    requirements_dict: [
+        requirement: {
+            "PO": [],
+            "safeguard": []
+        }
+    ]
+    """
+    candidate_safeguards = []
+    po_list = [po.lower().rstrip() for po in PO if po]
+    if "young users" in po_list and len(po_list) == 1:
+        return []
+    candidate_safeguards = safeguard_vector_store.get(where={"po": {"$in": po_list}})
+    safeguard_dict, safeguard_content = {}, []
+    for id, content, metadata in zip(candidate_safeguards['ids'], candidate_safeguards['documents'], candidate_safeguards['metadatas']):
+        safeguard_dict[content] = {
+            "metadata": metadata,
+            "rank": [],
+            "rrf_score": 0
+        }
+        safeguard_content.append(content)
+     
+    # Reranker
+    if using_reranker:
+        content_pairs, reranking_rank, reranking_results = [], [], []
+        for safeguard in safeguard_content:
+            content_pairs.append([requirement, safeguard])
+        safeguard_rerank_scores = reranker_model.compute_score(content_pairs)
+        for content_pair, score in zip(content_pairs, safeguard_rerank_scores):
+            reranking_rank.append((content_pair[1], score))
+        reranking_results = sorted(reranking_rank, key=lambda x: x[1], reverse=True)
+        for safeguard, score in reranking_results:
+            safeguard_dict[safeguard]['rank'].append(reranking_results.index((safeguard, score)) + 1)
+
+    #  BM25
+    if using_BM25:
+        with contextlib.redirect_stdout(io.StringIO()):
+            bm25_retriever = BM25(corpus=safeguard_content)
+            bm25_retriever.index(tokenize(safeguard_content))
+            bm25_results, scores = bm25_retriever.retrieve(tokenize(requirement), k = len(safeguard_content))
+            bm25_retrieval_rank = 1
+            for safeguard in bm25_results[0]:
+                safeguard_dict[safeguard]['rank'].append(bm25_retrieval_rank)
+                bm25_retrieval_rank += 1
+
+        # chroma retrieval
+    if using_chroma:
+        retrieved_safeguards = safeguard_vector_store.similarity_search_with_score(query=requirement, k=len(candidate_safeguards['ids']), filter={"po": {"$in": po_list}})
+        retrieval_rank = 1
+        for safeguard in retrieved_safeguards:
+            safeguard_dict[safeguard[0].page_content]['rank'].append(retrieval_rank)
+            retrieval_rank += 1
+
+    final_result = []
+    for safeguard in safeguard_content:
+        safeguard_dict[safeguard]['rrf_score'] = rrf(safeguard_dict[safeguard]['rank'])
+        final_result.append((safeguard_dict[safeguard]['rrf_score'], safeguard_dict[safeguard]['metadata']['safeguard_number'], safeguard, safeguard_dict[safeguard]['metadata']['po']))
+    final_result.sort(key=lambda x: x[0], reverse=True)
+        
+        # top k 
+    topk_final_result = final_result[:k]
+
+    return topk_final_result
+
+if __name__=="__main__":
+    embeddingmodel = TextEmb3LargeEmbedding(max_qpm=58)
+    embedding = EmbeddingFunction(embeddingmodel)
+    safeguard_vector_store = get_or_create_vector_base('safeguard_database', embedding)
+    reranker_model = FlagReranker(
+        '/root/PTR-LLM/tasks/pcf/model/bge-reranker-v2-m3',
+        use_fp16=True,
+        devices=["cpu"],   
+    )
+    requirement = """
+    Data Minimization Consent for incompatible purposes: Require consent for additional use of personal information not reasonably necessary to or incompatible with original purpose disclosure.
+    """
+    PO = ["Data Minimization & Purpose Limitation", "Transparency"]
+    final_result = retriever(
+            requirement,
+            PO,
+            safeguard_vector_store,
+            reranker_model,
+            using_reranker=True,
+            using_BM25=False,
+            using_chroma=True,
+            k=10
+            )
     print(final_result)