Update interface.py

interface.py

@@ -1,186 +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__":
+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)