| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| # 加载预训练的模型 | |
| tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base") | |
| model = AutoModelForSequenceClassification.from_pretrained("microsoft/codebert-base") | |
| # 定义漏洞检测函数 | |
| def detect_vulnerability(code): | |
| inputs = tokenizer(code, return_tensors="pt") | |
| outputs = model(**inputs) | |
| prediction = outputs.logits.argmax(-1).item() | |
| return "Buffer Overflow Detected" if prediction == 1 else "No Vulnerability Detected" | |
| # 创建 Gradio 界面 | |
| iface = gr.Interface( | |
| fn=detect_vulnerability, | |
| inputs="textbox", | |
| outputs="text", | |
| title="Buffer Overflow Detection", | |
| description="输入代码片段以检测是否存在缓冲区溢出漏洞。" | |
| ) | |
| iface.launch() | |