sangge/BackDoorBuster
sangge/BackDoorBuster
1
0
Fork 0
Code Issues Pull Requests Projects 1 Activity

fix:增加llm结果鲁棒性

Browse Source
This commit is contained in:
tritium0041 2024-06-04 17:13:30 +08:00
parent 65820966df
commit 72901463c6
1 changed files with 9 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
detection/GPTdetection.py
View File

@ -25,7 +25,7 @@ def detectGPT(content: str):
signal.signal(signal.SIGTERM, timeout_handler) signal.signal(signal.SIGTERM, timeout_handler)
signal.alarm(10) signal.alarm(10)
client = openai.OpenAI(base_url="https://api.xiaoai.plus/v1",api_key=api_key) client = openai.OpenAI(base_url="https://api.xiaoai.plus/v1", api_key=api_key)
text = content text = content
# client = openai.OpenAI(api_key="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") #测试用key # client = openai.OpenAI(api_key="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") #测试用key
response = client.chat.completions.create( response = client.chat.completions.create(
@ -60,7 +60,10 @@ def detectGPT(content: str):
classified_results = {"high": [], "medium": [], "low": [], "none": []} classified_results = {"high": [], "medium": [], "low": [], "none": []}
for res in res_json: for res in res_json:
try:
classified_results[res["Risk"]].append( classified_results[res["Risk"]].append(
(res["Line"], text.split("\n")[res["Line"] - 1].strip()) (res["Line"], text.split("\n")[res["Line"] - 1].strip())
) )
except IndexError:
pass
return classified_results return classified_results