Ddavid2002 commited on
Commit
938d8cb
1 Parent(s): 1435a9e

Add Gradio app for buffer overflow detection

Browse files
Files changed (2) hide show
  1. app.py +24 -0
  2. requirements.txt +3 -0
app.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoTokenizer, AutoModelForSequenceClassification
3
+
4
+ # 加载预训练的模型
5
+ tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")
6
+ model = AutoModelForSequenceClassification.from_pretrained("microsoft/codebert-base")
7
+
8
+ # 定义漏洞检测函数
9
+ def detect_vulnerability(code):
10
+ inputs = tokenizer(code, return_tensors="pt")
11
+ outputs = model(**inputs)
12
+ prediction = outputs.logits.argmax(-1).item()
13
+ return "Buffer Overflow Detected" if prediction == 1 else "No Vulnerability Detected"
14
+
15
+ # 创建 Gradio 界面
16
+ iface = gr.Interface(
17
+ fn=detect_vulnerability,
18
+ inputs="textbox",
19
+ outputs="text",
20
+ title="Buffer Overflow Detection",
21
+ description="输入代码片段以检测是否存在缓冲区溢出漏洞。"
22
+ )
23
+
24
+ iface.launch()
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ transformers
2
+ torch
3
+ gradio