Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+store/requirement_full_database/chroma.sqlite3 filter=lfs diff=lfs merge=lfs -text
+store/requirement_v1_database/chroma.sqlite3 filter=lfs diff=lfs merge=lfs -text
+store/requirement_v2_database/chroma.sqlite3 filter=lfs diff=lfs merge=lfs -text
+store/safeguard_database/chroma.sqlite3 filter=lfs diff=lfs merge=lfs -text

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

PROMPT_TEMPLATE.py

@@ -0,0 +1,29 @@
+former_string = """# Role
+## You are an expert in the field of law, and you are good at explaining why law requirements are related to their matching safeguards.
+
+# Task
+You need to analyze **requirement** and **privacy objective dict** provided by the user, each key in the **privacy objective dict** is a specific **privacy objective** and has corresponding **safeguards list**, you need to explain why each **safeguard** is related to the **requirement**. 
+
+# Output format
+For each **safeguard** in the **safeguards list**, explain its association with the requirement in the following format:
+{
+    "privacy objective":
+        [
+            {
+                "safeguard number": "xxx",
+                "safeguard description": "xxx",
+                "analysis": "xxx"
+            },
+            ...
+        ]
+}
+Please return your answers in JSON format."""
+input_format = """
+# Input
+Requirement:
+{requirement}
+Safeguards list:
+{safeguards}
+"""
+def prompt_template(requirement, safeguards):
+    return former_string + input_format.format(requirement=requirement, safeguards=safeguards)

README.md

@@ -1,12 +1,6 @@
 ---
-title: NLL Interface
-emoji: 🏆
-colorFrom: indigo
-colorTo: blue
+title: NLL_Interface
+app_file: interface.py
 sdk: gradio
-sdk_version: 5.23.3
-app_file: app.py
-pinned: false
+sdk_version: 5.23.2
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

conf/config.ini

@@ -0,0 +1,22 @@
+[DouBao128ProModelInfo]
+ACCESS_KEY = AKLTYjI0OWNiMGVmZGEwNDNhYjk3YzJhNDdlYTI1NTA5M2M
+SECRET_KEY = TldReVlUaGlNVE0wT0dVeE5ESTFOV0l3T1RKa1lXSm1aak0zTXpJeE5qVQ==
+BASE_URL = https://ark.cn-beijing.volces.com/api/v3
+API_KEY = 0c654012-8989-455f-8a5d-032fc067fbc8
+ENDPOINT_ID = ep-20241223113321-g47rr
+CYCLE_TIMES = 3
+MERGER_RETRY_TIMES = 4
+MAX_RETRY_TIMES = 4
+MAX_THREAD_NUM = 2
+
+
+[DouBaoPreviewModelInfo]
+ACCESS_KEY = AKLTYjI0OWNiMGVmZGEwNDNhYjk3YzJhNDdlYTI1NTA5M2M
+SECRET_KEY = TldReVlUaGlNVE0wT0dVeE5ESTFOV0l3T1RKa1lXSm1aak0zTXpJeE5qVQ==
+BASE_URL = https://ark.cn-beijing.volces.com/api/v3
+API_KEY = 0c654012-8989-455f-8a5d-032fc067fbc8
+ENDPOINT_ID = ep-20240923111539-mbwqc
+CYCLE_TIMES = 3
+MERGER_RETRY_TIMES = 4
+MAX_RETRY_TIMES = 4
+MAX_THREAD_NUM = 2

conf/logs.ini

@@ -0,0 +1,28 @@
+[loggers]
+keys = root,Robot
+
+[handlers]
+keys = consoleHandler
+
+[formatters]
+keys = simpleFormatter
+
+[logger_root]
+level = INFO
+handlers= consoleHandler
+
+[logger_Robot]
+level= INFO
+handlers = consoleHandler
+qualname = Robot
+propagate=0
+
+[handler_consoleHandler]
+class = StreamHandler
+level = INFO
+formatter = simpleFormatter
+
+
+[formatter_simpleFormatter]
+format = %(asctime)s %(levelname)s %(filename)s %(lineno)d %(message)s
+datefmt = %Y-%m-%d %H:%M:%S

doubao_service.py

@@ -0,0 +1,166 @@
+import time
+import requests
+import json
+from volcenginesdkarkruntime import Ark
+from util.config_util import read_config as config
+from util import logger
+import volcenginesdkcore
+import volcenginesdkark
+from volcenginesdkcore.rest import ApiException
+from util.logger_util import log_decorate
+
+
+class DouBaoService:
+
+    def __init__(self, model_name):
+        self.conf = config()[f"{model_name}ModelInfo"]
+        self.client = self.init_client()
+        self._complete_args = {}
+
+
+
+    def init_client(self):
+        base_url = self.conf["BASE_URL"]
+        ak = self.conf["ACCESS_KEY"]
+        sk = self.conf["SECRET_KEY"]
+        # api_key = self.conf["API_KEY"]
+        client = Ark(ak=ak, sk=sk, base_url=base_url)
+        # client = Ark(ak=api_key, base_url=base_url)
+        return client
+
+    def get_api_key(self):
+        configuration = volcenginesdkcore.Configuration()
+        configuration.ak = self.conf["ACCESS_KEY"]
+        configuration.sk = self.conf["SECRET_KEY"]
+        configuration.region = "cn-beijing"
+        endpoint_id = self.conf["ENDPOINT_ID"]
+
+        volcenginesdkcore.Configuration.set_default(configuration)
+
+        # use global default configuration
+        api_instance = volcenginesdkark.ARKApi()
+        get_api_key_request = volcenginesdkark.GetApiKeyRequest(
+            duration_seconds=30 * 24 * 3600,
+            resource_type="endpoint",
+            resource_ids=[
+                endpoint_id
+            ],
+        )
+
+        try:
+            resp = api_instance.get_api_key(get_api_key_request)
+            return resp.api_key
+        except ApiException as e:
+            logger.error(f"Exception when calling api: {e}")
+
+    def set_complete_args(self, temperature=None, top_p=None, max_token=None):
+        if temperature is not None:
+            self._complete_args["temperature"] = temperature
+        if top_p is not None:
+            self._complete_args["top_p"] = top_p
+        if max_token is not None:
+            self._complete_args["max_tokens"] = max_token
+
+    def form_user_role(self, content):
+        return {"role": "user", "content": content}
+
+    def form_sys_role(self, content):
+        return {"role": "system", "content": content}
+
+    def form_assistant_role(self, content):
+        return {"role": "assistant", "content": content}
+
+    @property
+    def complete_args(self):
+        return {"temperature": 0.01, "top_p": 0.7}
+
+    @log_decorate
+    def chat_complete(self, messages):
+
+        endpoint_id = self.conf["ENDPOINT_ID"]
+        completion = self.client.chat.completions.create(
+            model=endpoint_id,
+            messages=messages,
+            **self.complete_args
+        )
+        logger.info(f"complete doubao task, id: {completion.id}")
+        return completion.choices[0].message.content
+
+    def prd_to_keypoint(self, prd_content):
+
+        role_desc = {"role": "system", "content": PRD2KP_SYS}
+
+        messages = [
+            role_desc,
+            {"role": "user", "content": prd_content}
+        ]
+        return self.chat_complete(messages)
+
+    def prd_to_cases(self, prd_content, case_language="Chinese"):
+
+        role_desc = {"role": "system", "content": PRD_CASE_SYS[case_language]}
+
+        messages = [
+            role_desc,
+            {"role": "user", "content": prd_content}
+        ]
+        return self.chat_complete(messages)
+
+    def keypoint_to_case(self, key_points):
+
+        role_desc = {"role": "system", "content": KP2CASE_SYS}
+
+        messages = [
+            role_desc,
+            {"role": "user", "content": key_points}
+        ]
+        return self.chat_complete(messages)
+
+    def case_merge_together(self, case_suits):
+
+        role_desc = {"role": "system", "content": CASE_AGG_SYS}
+
+        content_case_suits = ""
+        for i, case_suit in enumerate(case_suits):
+            case_suit_expr = json.dumps(case_suit, indent=4, ensure_ascii=False)
+            content_case_suits += f"来自初级测试工程师{i + 1}的测试用例：\n```json\n{case_suit_expr}\n```\n"
+        messages = [
+            role_desc,
+            {"role": "user", "content": content_case_suits}
+        ]
+        completion = self.chat_complete(messages)
+        return completion
+
+    def cycle_more_case(self, prd_content, case_language="Chinese"):
+
+        role_desc = {"role": "system", "content": PRD_CASE_SYS[case_language]}
+
+        messages = [
+            role_desc,
+            {"role": "user", "content": PRD_CASE_1[case_language] + prd_content + "\n" + PRD_CASE_2[case_language]}
+        ]
+
+        result = []
+
+        for sys in MORE_CASE_PROMPT[case_language]:
+            if sys:
+                messages.append({"role": "user", "content": sys})
+            reply = self.chat_complete(messages)
+            result.append(reply)
+            messages.append({"role": "assistant", "content": reply})
+            time.sleep(10)
+        return result
+
+
+if __name__ == "__main__":
+    cli = DouBaoService("DouBao128Pro")
+    # print(cli.get_api_key())
+    # prd_content = requests.get("https://tosv.byted.org/obj/music-qa-bucket/xmind-test/de3ebc67410c43603034e21bfefa76a0.md").text
+    # aa = cli.cycle_more_case(prd_content, "English")
+    # print(aa)
+
+    print(cli.chat_complete(messages=[
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "Introduce LLM shortly."},
+    ]))
+

interface.py

@@ -0,0 +1,186 @@
+import pandas as pd
+import json
+import re
+from json import loads, JSONDecodeError
+import sys
+import os
+import ast
+from util.vector_base import EmbeddingFunction, get_or_create_vector_base
+from doubao_service import DouBaoService
+from PROMPT_TEMPLATE import prompt_template
+from util.Embeddings import TextEmb3LargeEmbedding
+from langchain_core.documents import Document
+from FlagEmbedding import FlagReranker
+from retriever import retriever
+import time
+from bm25s import BM25, tokenize
+import contextlib
+import io
+
+import gradio as gr
+import time
+
+client = DouBaoService("DouBao128Pro")
+embeddingmodel = TextEmb3LargeEmbedding(max_qpm=58)
+embedding = EmbeddingFunction(embeddingmodel)
+safeguard_vector_store = get_or_create_vector_base('safeguard_database', embedding)
+
+# reranker_model = FlagReranker(
+#     'C://Users//Admin//Desktop//PDPO//NLL_LLM//model//bge-reranker-v2-m3',
+#     use_fp16=True,
+#     devices=["cpu"],   
+# )
+
+OPTIONS = ['AI Governance',
+ 'Data Accuracy',
+ 'Data Minimization & Purpose Limitation',
+ 'Data Retention',
+ 'Data Security',
+ 'Data Sharing',
+ 'Individual Rights',
+ 'Privacy by Design',
+ 'Transparency']
+
+
+def format_model_output(raw_output):
+    """
+    处理模型输出：
+    - 将 \n 转换为实际换行
+    - 提取 ```json ``` 中的内容并格式化为可折叠的 JSON
+    """
+    formatted = raw_output.replace('\\n', '\n')  
+    def replace_json(match):
+        json_str = match.group(1).strip()
+        try:
+            json_obj = loads(json_str)
+            return f"```json\n{json.dumps(json_obj, indent=2, ensure_ascii=False)}\n```"
+        except JSONDecodeError:
+            return match.group(0) 
+  
+    formatted = re.sub(r'```json\n?(.*?)\n?```', replace_json, formatted, flags=re.DOTALL)
+    return ast.literal_eval(formatted) 
+
+def model_predict(input_text, if_split_po, topk, selected_items):
+    """
+    selected_items: 用户选择的项目（可能是["All"]或具体PO）
+    """
+    requirement = input_text
+    requirement = requirement.replace("\t", "").replace("\n", "").replace("\r", "")
+    if "All" in selected_items:
+        PO = OPTIONS
+    else:
+        PO = selected_items
+    if topk:
+        topk = int(topk)
+    else:
+        topk = 10
+    final_result = retriever(
+                requirement,
+                PO,
+                safeguard_vector_store,
+                reranker_model=None,
+                using_reranker=False,
+                using_BM25=False,
+                using_chroma=True,
+                k=topk,
+                if_split_po=if_split_po
+            )
+    mapping_safeguards = {}
+    for safeguard in final_result:
+        if safeguard[3] not in mapping_safeguards:
+            mapping_safeguards[safeguard[3]] = []
+        mapping_safeguards[safeguard[3]].append(
+            {
+                "Score": safeguard[0],
+                "Safeguard Number": safeguard[1],
+                "Safeguard Description": safeguard[2]
+            }
+        )
+    prompt = prompt_template(requirement, mapping_safeguards)
+    response = client.chat_complete(messages=[
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": prompt},
+    ])
+    # return {"requirement": requirement, "safeguards": mapping_safeguards}
+    print("requirement:", requirement)
+    print("mapping safeguards:", mapping_safeguards)
+    print("response:", response)
+    return {"requirement": requirement, "safeguards": format_model_output(response)}
+
+with gr.Blocks(title="New Law Landing") as demo:
+    gr.Markdown("## 🏙️ New Law Landing")
+    
+    requirement = gr.Textbox(label="Input Requirements", placeholder="Example: Data Minimization Consent for incompatible purposes")
+    details = gr.Textbox(label="Input Details", placeholder="Example: Require consent for...")
+    
+    # 修改为 Number 输入组件
+    topk = gr.Number(
+        label="Top K safeguards", 
+        value=10,  
+        precision=0,  
+        minimum=1,  
+        interactive=True
+    )
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            if_split_po = gr.Checkbox(
+                label="If Split Privacy Objective", 
+                value=True, 
+                info="Recall K Safeguards for each Privacy Objective"
+            )
+        with gr.Column(scale=1):
+            all_checkbox = gr.Checkbox(
+                label="ALL Privacy Objective", 
+                value=True, 
+                info="No specific Privacy Objective is specified"
+            )
+        with gr.Column(scale=4):
+            PO_checklist = gr.CheckboxGroup(
+                label="Choose Privacy Objective",
+                choices=OPTIONS,
+                value=[],  
+                interactive=True
+            )
+
+    submit_btn = gr.Button("Submit", variant="primary")
+    result_output = gr.JSON(label="Related safeguards", open=True)
+    
+
+    def sync_checkboxes(selected_items, all_selected):
+        if len(selected_items) > 0:
+            return False
+        return all_selected
+    
+    PO_checklist.change(
+        fn=sync_checkboxes,
+        inputs=[PO_checklist, all_checkbox],
+        outputs=all_checkbox
+    )
+    
+    def sync_all(selected_all, current_selection):
+        if selected_all:
+            return []
+        return current_selection
+    
+    all_checkbox.change(
+        fn=sync_all,
+        inputs=[all_checkbox, PO_checklist],
+        outputs=PO_checklist
+    )
+ 
+    def process_inputs(requirement, details, topk, if_split_po, all_selected, PO_selected):
+        input_text = requirement + ": " + details
+        if all_selected:
+            return model_predict(input_text, if_split_po, int(topk), ["All"])
+        else:
+            return model_predict(input_text, if_split_po, int(topk), PO_selected)
+    
+    submit_btn.click(
+        fn=process_inputs,
+        inputs=[requirement, details, topk, if_split_po, all_checkbox, PO_checklist],
+        outputs=[result_output]
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True)

retriever.py

@@ -0,0 +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
+            )
+    print(final_result)

store/requirement_full_database/8879b034-d26b-4dd9-bdc6-9a0751f8eeeb/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b8d4b3825a7c7a773e22fa3eeef0e7d15a695f5c4183aeff5beb07741a68679
+size 12428000

store/requirement_full_database/8879b034-d26b-4dd9-bdc6-9a0751f8eeeb/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8a3ec48846fc6fdfaef19f5ed2508f0bf3da4a3c93b0f6b3dd21f0a22ec1026
+size 100

store/requirement_full_database/8879b034-d26b-4dd9-bdc6-9a0751f8eeeb/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fc19b1997119425765295aeab72d76faa6927d4f83985d328c26f20468d6cc76
+size 4000

store/requirement_full_database/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:debbda97fa444d1beed59205da5310caa07fc5e1fea98ee5d217bb1cd86b3312
+size 2031616

store/requirement_v1_database/6db99751-9b95-42b7-ae30-46ba43f95c27/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b8d4b3825a7c7a773e22fa3eeef0e7d15a695f5c4183aeff5beb07741a68679
+size 12428000

store/requirement_v1_database/6db99751-9b95-42b7-ae30-46ba43f95c27/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8a3ec48846fc6fdfaef19f5ed2508f0bf3da4a3c93b0f6b3dd21f0a22ec1026
+size 100

store/requirement_v1_database/6db99751-9b95-42b7-ae30-46ba43f95c27/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e0c0e05e944e59611aa54c4b0b708f835a8a0daf4baec11208e3f60773b22d89
+size 4000

store/requirement_v1_database/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f177f30d3354a5cf2c14378bfe1b73755684d9c5fc38046f9e4339a1180af0a2
+size 1212416

store/requirement_v2_database/c1b4f057-aa88-49ff-a2ac-08fc9d60804c/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b8d4b3825a7c7a773e22fa3eeef0e7d15a695f5c4183aeff5beb07741a68679
+size 12428000

store/requirement_v2_database/c1b4f057-aa88-49ff-a2ac-08fc9d60804c/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8a3ec48846fc6fdfaef19f5ed2508f0bf3da4a3c93b0f6b3dd21f0a22ec1026
+size 100

store/requirement_v2_database/c1b4f057-aa88-49ff-a2ac-08fc9d60804c/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9a911716b3f8450b156db9e04a3c81548395cabac0c846dd1e1eef832991c120
+size 4000

store/requirement_v2_database/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b6d2f15577b5bcd6c8ce845d894cc17196e71c71c3909eec73f9bfedee400c90
+size 991232

store/safeguard_database/1ae9d702-e220-41de-95e3-e603f3a12409/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b8d4b3825a7c7a773e22fa3eeef0e7d15a695f5c4183aeff5beb07741a68679
+size 12428000

store/safeguard_database/1ae9d702-e220-41de-95e3-e603f3a12409/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8a3ec48846fc6fdfaef19f5ed2508f0bf3da4a3c93b0f6b3dd21f0a22ec1026
+size 100

store/safeguard_database/1ae9d702-e220-41de-95e3-e603f3a12409/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c710099c23f51c095ec65b43539cd7534bb745aebf55fef0a9c06229121caca
+size 4000

store/safeguard_database/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0730981819c475f93d173a347f891b5e28e02cce24a63a578ca43c98a13d9c10
+size 5095424

util/Embeddings.py

@@ -0,0 +1,195 @@
+import os
+from copy import copy
+from typing import Dict, List, Optional, Tuple, Union
+import numpy as np
+import time
+from functools import wraps
+
+os.environ['CURL_CA_BUNDLE'] = ''
+from dotenv import load_dotenv, find_dotenv
+_ = load_dotenv(find_dotenv())
+
+
+class BaseEmbeddings:
+    """
+    Base class for embeddings
+    """
+    def __init__(self, path: str, is_api: bool) -> None:
+        self.path = path
+        self.is_api = is_api
+
+    def get_embedding(self, text: str, model: str) -> List[float]:
+        raise NotImplementedError
+
+    @classmethod
+    def cosine_similarity(cls, vector1: List[float], vector2: List[float]) -> float:
+        """
+        calculate cosine similarity between two vectors
+        """
+        dot_product = np.dot(vector1, vector2)
+        magnitude = np.linalg.norm(vector1) * np.linalg.norm(vector2)
+        if not magnitude:
+            return 0
+        return dot_product / magnitude
+
+
+class OpenAIEmbedding(BaseEmbeddings):
+    """
+    class for OpenAI embeddings
+    """
+    def __init__(self, path: str = '', is_api: bool = True) -> None:
+        super().__init__(path, is_api)
+        if self.is_api:
+            from openai import OpenAI
+            self.client = OpenAI()
+            self.client.api_key = os.getenv("OPENAI_API_KEY")
+            self.client.base_url = os.getenv("OPENAI_BASE_URL")
+
+    def get_embedding(self, text: str, model: str = "text-embedding-3-large") -> List[float]:
+        if self.is_api:
+            text = text.replace("\n", " ")
+            return self.client.embeddings.create(input=[text], model=model).data[0].embedding
+        else:
+            raise NotImplementedError
+
+class JinaEmbedding(BaseEmbeddings):
+    """
+    class for Jina embeddings
+    """
+    def __init__(self, path: str = 'jinaai/jina-embeddings-v2-base-zh', is_api: bool = False) -> None:
+        super().__init__(path, is_api)
+        self._model = self.load_model()
+
+    def get_embedding(self, text: str) -> List[float]:
+        return self._model.encode([text])[0].tolist()
+
+    def load_model(self):
+        import torch
+        from transformers import AutoModel
+        if torch.cuda.is_available():
+            device = torch.device("cuda")
+        else:
+            device = torch.device("cpu")
+        model = AutoModel.from_pretrained(self.path, trust_remote_code=True).to(device)
+        return model
+
+class ZhipuEmbedding(BaseEmbeddings):
+    """
+    class for Zhipu embeddings
+    """ 
+    def __init__(self, path: str = '', is_api: bool = True) -> None:
+        super().__init__(path, is_api)
+        if self.is_api:
+            from zhipuai import ZhipuAI
+            self.client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"))
+
+    def get_embedding(self, text: str) -> List[float]:
+        response = self.client.embeddings.create(
+        model="embedding-2",
+        input=text,
+        )
+        return response.data[0].embedding
+    
+
+class DashscopeEmbedding(BaseEmbeddings):
+    """
+    class for Dashscope embeddings
+    """
+    def __init__(self, path: str = '', is_api: bool = True) -> None:
+        super().__init__(path, is_api)
+        if self.is_api:
+            import dashscope
+            dashscope.api_key = os.getenv("DASHSCOPE_API_KEY")
+            self.client = dashscope.TextEmbedding
+
+    def get_embedding(self, text: str, model: str='text-embedding-v1') -> List[float]:
+        response = self.client.call(
+            model=model,
+            input=text
+        )
+        return response.output['embeddings'][0]['embedding']
+
+
+class BgeEmbedding(BaseEmbeddings):
+    """
+    class for BGE embeddings
+    """
+
+    def __init__(self, path: str = 'BAAI/bge-en-icl', is_api: bool = False) -> None:
+        super().__init__(path, is_api)
+        self._model, self._tokenizer = self.load_model(path)
+
+    def get_embedding(self, text: str) -> List[float]:
+        import torch
+        encoded_input = self._tokenizer([text], padding=True, truncation=True, return_tensors='pt')
+        encoded_input = {k: v.to(self._model.device) for k, v in encoded_input.items()}
+        with torch.no_grad():
+            model_output = self._model(**encoded_input)
+            sentence_embeddings = model_output[0][:, 0]
+        sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
+        return sentence_embeddings[0].tolist()
+
+    def load_model(self, path: str):
+        import torch
+        from transformers import AutoModel, AutoTokenizer
+        if torch.cuda.is_available():
+            device = torch.device("cuda")
+        else:
+            device = torch.device("cpu")
+        tokenizer = AutoTokenizer.from_pretrained(path)
+        model = AutoModel.from_pretrained(path).to(device)
+        model.eval()
+        return model, tokenizer
+
+
+def rate_limiter():
+    def rate_limiter_decorator(func):
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            max_calls_per_minute = self.max_qpm 
+            interval = 60 / max_calls_per_minute
+            current_time = time.time()
+            
+            # Check if there's a record of the last call, if not set it to 0
+            if not hasattr(self, '_last_called'):
+                self._last_called = 0
+            elapsed_time = current_time - self._last_called
+            if elapsed_time < interval:
+                time_to_wait = interval - elapsed_time
+                if self.silent is False:
+                    print(f"## Rate limit reached. Waiting for {time_to_wait:.2f} seconds.")
+                time.sleep(time_to_wait)
+            result = func(self, *args, **kwargs)
+            self._last_called = time.time()  
+            return result
+        return wrapper
+
+    return rate_limiter_decorator
+
+
+class TextEmb3LargeEmbedding(BaseEmbeddings):
+    """
+    class for text-embedding-3-large embeddings
+    """ 
+    def __init__(self, max_qpm, is_silent=False):
+        from langchain_openai import AzureOpenAIEmbeddings
+
+        ## https://gpt.bytedance.net/gpt_openapi/
+        base_url = "https://search-va.byteintl.net/gpt/openapi/online/v2/crawl"
+        api_version = "2024-03-01-preview"
+        ak = "5dXdIKxZc8JWVVgvX0DN92HWIYb9NfEb_GPT_AK"
+        model_name = "text-embedding-3-large"
+        api_type = "azure"       
+        self.llm = AzureOpenAIEmbeddings(
+            azure_endpoint=base_url,
+            openai_api_version=api_version,
+            deployment=model_name,
+            openai_api_key=ak,
+            openai_api_type=api_type,
+        )
+        self.max_qpm = max_qpm
+        self.silent = is_silent
+
+    @rate_limiter()
+    def get_embedding(self, text: str):
+        return self.llm.embed_query(text)

util/__init__.py

@@ -0,0 +1,3 @@
+from .logger_util import get_logger
+
+logger = get_logger()

util/config_util.py

@@ -0,0 +1,21 @@
+import os
+import json
+import logging
+import bytedenv
+import configparser
+
+ROOT_PATH = os.path.split(os.path.split(os.path.abspath(__file__))[0])[0]
+
+
+def read_config():
+    config_file = ROOT_PATH + "\conf\config.ini"
+    config_ini = configparser.ConfigParser()
+    config_ini.read(config_file)
+    model_name = "DouBao128Pro"
+    return config_ini
+
+def read_json(filepath):
+    with open(filepath, 'r') as f:
+        result = json.load(f)
+        return result
+

util/logger_util.py

@@ -0,0 +1,23 @@
+import os
+import logging
+import logging.config
+import traceback
+from functools import wraps
+
+
+def get_logger():
+    root_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    logging.config.fileConfig(os.path.join(root_path, "conf", "logs.ini"))
+    logger = logging.getLogger("Robot")
+    return logger
+
+
+def log_decorate(func):
+    @wraps(func)
+    def log(*args, **kwargs):
+        logger = get_logger()
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            logger.error(f"{func.__name__} is error, logId: {e.args},  errMsg is: {traceback.format_exc()}")
+    return log

util/vector_base.py

@@ -0,0 +1,79 @@
+import sys
+from langchain_chroma import Chroma
+from langchain_core.documents import Document
+sys.path.append('C://Users//Admin//Desktop//PDPO//NLL_LLM//util')
+from Embeddings import TextEmb3LargeEmbedding
+from pathlib import Path
+import time
+
+class EmbeddingFunction():
+    def __init__(self, embeddingmodel):
+        self.embeddingmodel = embeddingmodel
+    def embed_query(self, query):
+        return list(self.embeddingmodel.get_embedding(query))
+    def embed_documents(self, documents):
+        return [self.embeddingmodel.get_embedding(document) for document in documents]
+
+def get_or_create_vector_base(collection_name: str, embedding, documents=None) -> Chroma:
+    """
+    判断vector store是否已经构建好，如果没有构建好，则先初始化vector store。不使用embed_documents
+    方法批量初始化vector store而是for循环逐个加入，同时使用sleep，以此避免调用openai的接口达到最大
+    上限而导致初始化失败。
+    """
+    persist_directory = "C://Users//Admin//Desktop//PDPO//NLL_LLM//store//" +collection_name
+    persist_path = Path(persist_directory)
+    if not persist_path.exists and not documents:
+        raise ValueError("vector store does not exist and documents is empty")
+    elif persist_path.exists():
+        print("vector store already exists")
+        vector_store = Chroma(
+            collection_name=collection_name,
+            embedding_function=embedding,
+            persist_directory=persist_directory
+        )
+    else:
+        print("start creating vector store")
+        vector_store = Chroma(
+            collection_name=collection_name,
+            embedding_function=embedding,
+            persist_directory=persist_directory
+        )
+        for document in documents:
+            vector_store.add_documents(documents=[document])
+            time.sleep(1)
+    return vector_store
+
+if __name__=="__main__":
+    import pandas as pd
+    requirements_data = pd.read_csv("/root/PTR-LLM/tasks/pcf/reference/NLL_DATA_NEW_Test.csv")
+    requirements_dict_v2 = {}
+    for index, row in requirements_data.iterrows():
+        requirement = row['Requirement'].split("- ")[1]
+        requirement = requirement + ": " + row['Details']
+        requirement = requirement.replace('\n', ' ').replace('\r', ' ').replace('\t', ' ')
+        if requirement not in requirements_dict_v2:
+            requirements_dict_v2[requirement] = {
+                'PO': set(),
+                'safeguard': set()
+            }
+        requirements_dict_v2[requirement]['PO'].add(row['PCF-Privacy Objective'].lower().rstrip() if isinstance(row['PCF-Privacy Objective'], str) else None)
+        requirements_dict_v2[requirement]['safeguard'].add(row['Safeguard'].lower().rstrip())
+    index = 0
+    documents = []
+    for key, value in requirements_dict_v2.items():
+        page_content = key
+        metadata = {
+            "index": index,
+            "version":2,
+            "PO": str([po for po in value['PO'] if po]),
+            "safeguard":str([safeguard for safeguard in value['safeguard']])
+        }
+        index += 1
+        document=Document(
+            page_content=page_content,
+            metadata=metadata
+        )
+        documents.append(document)
+    embeddingmodel = TextEmb3LargeEmbedding(max_qpm=58)
+    embedding = EmbeddingFunction(embeddingmodel)
+    requirement_v2_vector_store = get_or_create_vector_base('requirement_v2_database', embedding, documents)