Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import warnings | |
| warnings.filterwarnings("ignore") | |
| import pandas as pd | |
| import numpy as np | |
| import faiss | |
| import ast | |
| import torch.nn.functional as F | |
| import torch | |
| from transformers import AutoModel, AutoTokenizer | |
| Encoding_model = 'jinaai/jina-embeddings-v2-base-zh' | |
| model = AutoModel.from_pretrained(Encoding_model, trust_remote_code=True, torch_dtype=torch.bfloat16) | |
| model#.to("cuda") | |
| similarity_model = 'Alibaba-NLP/gte-multilingual-base' | |
| similarity_tokenizer = AutoTokenizer.from_pretrained(similarity_model) | |
| similarity_model = AutoModel.from_pretrained(similarity_model, trust_remote_code=True)#.to("cuda") | |
| def get_not_empty_data(df,x_column="text",y_column="label"): | |
| df = df[df[y_column] != "[]"].reset_index(drop=True) | |
| res_dict = {} | |
| for idx in df.index: | |
| if df.loc[idx,x_column] not in res_dict: | |
| res_dict[df.loc[idx,x_column]] = ast.literal_eval(df.loc[idx,y_column]) | |
| else: | |
| res_dict[df.loc[idx,x_column]] += ast.literal_eval(df.loc[idx,y_column]) | |
| res_dict = {k:list(set(v)) for k,v in res_dict.items()} | |
| df_dict = pd.DataFrame({"x":res_dict.keys(),"y":res_dict.values()}) | |
| return df_dict | |
| data_all = pd.read_excel("data_Excel_format.xlsx") | |
| df_dict_all = get_not_empty_data(data_all) | |
| x_dict = df_dict_all["x"].values | |
| y_dict = df_dict_all["y"].values | |
| def calc_scores(x): | |
| return (x[:1] @ x[1:].T) | |
| def get_idxs(threshold,max_len,arr): | |
| res = np.where(arr >= threshold)[0] | |
| if len(res)<max_len: | |
| return res | |
| res = res[np.argsort(-arr[res])][:3] | |
| return res | |
| def merge_set_to_list(set_list): | |
| res = set() | |
| for i in set_list: | |
| res = res | i | |
| return res | |
| def get_predict_result(index,score,threshold,max_len): | |
| score = score.flatten() | |
| index = index.flatten() | |
| index_of_index = np.where(score >= threshold)[0] | |
| if len(index_of_index)>=max_len: | |
| index_of_index = index_of_index[np.argsort(-index[index_of_index])][:3] | |
| if len(index_of_index)==0: | |
| return {},[] | |
| res_index = index[index_of_index] | |
| res = merge_set_to_list([set(i) for i in y_dict[res_index]]) | |
| return res,x_dict[res_index] | |
| vec = np.empty(shape=[0,768],dtype="float32") | |
| bsize = 256 | |
| with torch.no_grad(): | |
| for i in range(0,len(x),bsize): | |
| tmp = model.encode(x[i:i+bsize]) | |
| vec = np.concatenate([vec,tmp]) | |
| index = faiss.IndexFlatIP(768) | |
| faiss.normalize_L2(vec) | |
| index.add(vec) | |
| faiss.write_index(index,"all_index.faiss") | |
| index = faiss.read_index("all_index.faiss") | |
| def predict_label(x,threshold=0.85,n_nearest=10,max_result_len=3): | |
| bsize=1 | |
| y_pred = [] | |
| with torch.no_grad(): | |
| for i in range(0,len(x),bsize): | |
| sentences = x[i:i+bsize] | |
| vec = model.encode(sentences) | |
| faiss.normalize_L2(vec) | |
| scores, indexes = index.search(vec,n_nearest) | |
| x_pred = np.array([[sentences[j]]+s.tolist() for j,s in enumerate(x_dict[indexes])]) | |
| batch_dict = similarity_tokenizer(x_pred.flatten().tolist(), max_length=768, padding=True, truncation=True, return_tensors='pt')#.to("cuda") | |
| outputs = similarity_model(**batch_dict) | |
| dimension=768 | |
| embeddings = outputs.last_hidden_state[:, 0][:dimension] | |
| embeddings = F.normalize(embeddings, p=2, dim=1) | |
| embeddings = embeddings.view(len(x_pred),n_nearest+1,dimension).detach().cpu().numpy() | |
| scores = [calc_scores(embeddings[b]) for b in range(embeddings.shape[0])] | |
| pred = [get_predict_result(indexes[k],scores[k],threshold=threshold,max_len=max_result_len) for k in range(len(scores))] | |
| y_pred.append([i[0] for i in pred]) | |
| return y_pred | |
| CSS_Content = """ | |
| <!DOCTYPE html> | |
| <html lang="en"> | |
| <head> | |
| <meta charset="UTF-8"> | |
| <meta name="viewport" content="width=device-width, initial-scale=1.0"> | |
| <style> | |
| #custom_id { | |
| border: 2px solid red; | |
| padding: 10px; | |
| background-color: lightgray; | |
| } | |
| </style> | |
| </head> | |
| </html> | |
| <span style="color: red;line-height:1;">红色字体：潜在风险</span><br> | |
| <span style="color: blue;line-height:1;">蓝色字体：权限获取</span><br> | |
| <span style="color: purple;line-height:1;">紫色字体：数据收集</span><br> | |
| <span style="color: green;line-height:1;">绿色字体：数据、权限管理</span><br> | |
| <span style="color: brown;line-height:1;">棕色字体：共享、委托、转让、公开（披露）</span><br> | |
| """ | |
| color_dict = {"潜在风险":"red", | |
| "权限获取":"blue", | |
| "数据收集":"purple", | |
| "数据、权限管理":"green", | |
| "共享、委托、转让、公开（披露）":"brown" | |
| } | |
| def generate_HTML(text,threshold=0.85,n_nearest=10,max_result_len=3): | |
| sentences = text.split("\n") | |
| sentences = [i for i in map(lambda x:x.split("。"),sentences)] | |
| res = CSS_Content | |
| for paragraph in sentences: | |
| tmp_res = [] | |
| pred_label = predict_label(paragraph,threshold,n_nearest,max_result_len) | |
| for i,x in enumerate(pred_label): | |
| pre = "<span" | |
| if len(x[0])>0: | |
| for j in color_dict.keys(): #color dict重要性递减，所以只取第一个标签的颜色 | |
| if j in x[0]: | |
| pre += f' style="color: {color_dict[j]};line-height:1;"' | |
| break | |
| tmp_res.append(pre+">"+paragraph[i]+"</span>") | |
| res += "。".join(tmp_res) | |
| res += "<br>" | |
| return res | |
| with gr.Blocks() as demo: | |
| with gr.Row(): | |
| input_text = gr.Textbox(lines=25,label="输入") | |
| with gr.Row(): | |
| threshold = gr.Slider(minimum=0.5,maximum=0.85,value=0.75,step=0.05,interactive=True,label="相似度阈值") | |
| n_nearest = gr.Slider(minimum=3,maximum=10,value=10,step=1,interactive=True,label="粗筛语句数量") | |
| max_result_len = gr.Slider(minimum=1,maximum=5,value=3,step=1,interactive=True,label="精筛语句数量") | |
| with gr.Row(): | |
| submit_button = gr.Button("检测") | |
| with gr.Row(): | |
| output_text = gr.HTML(CSS_Content) | |
| output_text.elem_id="custom_id" | |
| submit_button.click(fn=generate_HTML, inputs=[input_text,threshold,n_nearest,max_result_len], outputs=output_text) | |
| demo.launch() |