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()