Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from .expert_malware_detector import expert_detect_malware | |
| from .discriminative_chatter_detector import DiscriminativeChatterDetector | |
| from .generative_malware_detector import GenerativeMalwareDetector | |
| detector_a = DiscriminativeChatterDetector(dataset_name='blackbriar') | |
| detector_c = GenerativeMalwareDetector() | |
| def detect_malware_a(code: str, dependencies: str) -> str: | |
| """Model A detection""" | |
| result = detector_a.predict(code) | |
| return f"**Prediction:** {result['label']}" | |
| def detect_malware_b(code: str, dependencies: str) -> str: | |
| """Model B detection""" | |
| result = expert_detect_malware(source_code=code, dependencies=dependencies) | |
| return f"**Prediction:** {result}" | |
| def detect_malware_c(code: str, dependencies: str) -> str: | |
| """Model C detection""" | |
| result = detector_c.detect(code, dependencies) | |
| return f"**Prediction:** {result['label']}" | |
| def clear_all_outputs(): | |
| """ | |
| Source - https://stackoverflow.com/a | |
| Posted by Vaishnav | |
| Retrieved 2025-11-17, License - CC BY-SA 4.0 | |
| """ | |
| return None, None, None, None, None | |
| # hwr_theme = gr.themes.Default(primary_hue=gr.themes.colors.rose, secondary_hue=gr.themes.colors.red) | |
| # with gr.Blocks(theme=hwr_theme) as demo: | |
| with gr.Blocks() as demo: | |
| gr.Markdown(""" | |
| # Malware Detection Demo | |
| Test three different AI models for detecting malicious code. | |
| One demo uses a rules-based model, one uses predictive ML ('discriminative') and one uses generative ML. | |
| Can you guess which is which? | |
| **Your task:** Try various inputs and track the responses to guess which demo is using which model type. | |
| """) | |
| with gr.Row(): | |
| code_text = gr.Textbox( | |
| label="Enter source code to analyze", | |
| placeholder="Add your code here...", | |
| lines=3 | |
| ) | |
| dependencies_text = gr.Textbox( | |
| label="Enter dependency specificiation to analyze", | |
| placeholder="Add your dependencies here...", | |
| lines=3 | |
| ) | |
| gr.Markdown("### Compare All Three Models") | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.Markdown("#### Model A") | |
| output_a = gr.Markdown() | |
| btn_a = gr.Button("Analyze with Model A", variant="primary") | |
| with gr.Column(): | |
| gr.Markdown("#### Model B") | |
| output_b = gr.Markdown() | |
| btn_b = gr.Button("Analyze with Model B", variant="primary") | |
| with gr.Column(): | |
| gr.Markdown("#### Model C") | |
| output_c = gr.Markdown() | |
| btn_c = gr.Button("Analyze with Model C", variant="primary") | |
| # Example inputs | |
| gr.Examples( | |
| examples=[ | |
| ["import os\nos.system('rm -rf /')", "requests==2.28.0\npandas==1.5.0"], | |
| ["print('Hello World')", "numpy==1.24.0"], | |
| ["from dspy import answer_every_question_no_hallucinations_trust_me", "dspy-ai==42.3.14"], | |
| ], | |
| inputs=[code_text, dependencies_text] | |
| ) | |
| clear_btn = gr.Button("Let's try again.") | |
| # Connect buttons to functions | |
| btn_a.click(fn=detect_malware_a, inputs=[code_text, dependencies_text], outputs=output_a) | |
| btn_b.click(fn=detect_malware_b, inputs=[code_text, dependencies_text], outputs=output_b) | |
| btn_c.click(fn=detect_malware_c, inputs=[code_text, dependencies_text], outputs=output_c) | |
| clear_btn.click(fn=clear_all_outputs, inputs=None, outputs=[code_text, dependencies_text, output_a, output_b, output_c]) | |