update app.py

app.py

@@ -1,3 +1,4 @@
+import textwrap
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -61,97 +62,97 @@ def fix_code(language, vulnerable_code):
     else:
         return result.strip()
 
-# UI Layout
-with gr.Blocks(title="PyJavaCPP Vuln-Fixer", theme=gr.themes.Base()) as demo:
-    gr.Markdown("# 🛡️ PyJavaCPP Vulnerability Fixer (CPU)")
-    gr.Markdown("Select your language, paste your code, and get a secured version of your code!")
-    
-    with gr.Row():
-        with gr.Column():
-            lang_input = gr.Dropdown(
-                choices=["python", "java", "cpp"], 
-                value="python", 
-                label="Target Language"
-            )
-            code_input = gr.Code(
-                label="Vulnerable Code", 
-                language="python", 
-                lines=12
-            )
-            submit_btn = gr.Button("Secure My Code ✨", variant="primary")
-        
-        with gr.Column():
-            code_output = gr.Code(
-                label="Fixed Code", 
-                language="python", 
-                lines=12, 
-                interactive=False
-            )
-
-    # Example Snippets for quick testing
-    gr.Examples(
-        examples=[
-            ["python", r"""import os
+EXAMPLES = [
+    [
+        "python",
+        textwrap.dedent("""\
+            import os
             from flask import Flask, request
-            
+
             app = Flask(__name__)
-            
+
             @app.route("/run")
             def run():
                 cmd = request.args.get("cmd")
-                # Vulnerable: Command Injection
-                
                 return os.popen(cmd).read()
-            
+
             if __name__ == "__main__":
-                app.run(debug=False)"""],
-            
-            ["java", r"""import java.sql.*;
+                app.run(debug=False)"""),
+    ],
+    [
+        "java",
+        textwrap.dedent("""\
+            import java.sql.*;
             import javax.servlet.http.*;
-            
+
             public class UserServlet extends HttpServlet {
                 public void doGet(HttpServletRequest req, HttpServletResponse res) {
                     try {
                         String id = req.getParameter("id");
                         Connection conn = DriverManager.getConnection("jdbc:mysql://localhost/db", "user", "pass");
                         Statement stmt = conn.createStatement();
-                        // Vulnerable: SQL Injection
                         ResultSet rs = stmt.executeQuery("SELECT * FROM users WHERE id='" + id + "'");
                     } catch (Exception e) {
                         e.printStackTrace();
                     }
                 }
-                
-                public static void main(String[] args) {
-                    System.out.println("Servlet loaded.");
-                }
-            }"""],
-            
-            ["cpp", r"""#include <iostream>
+            }"""),
+    ],
+    [
+        "cpp",
+        textwrap.dedent("""\
+            #include <iostream>
             #include <cstring>
-            
+
             void login(char *input) {
                 char password[8];
-                // Vulnerable: Buffer Overflow
                 strcpy(password, input);
             }
-            
+
             int main(int argc, char *argv[]) {
                 if (argc > 1) {
                     login(argv[1]);
                 }
-                
                 return 0;
-            }"""]
-        ],
-        inputs=[lang_input, code_input]
+            }"""),
+    ],
+]
+
+# UI Layout
+with gr.Blocks(title="PyJavaCPP Vuln-Fixer") as demo:
+    gr.Markdown("# 🛡️ PyJavaCPP Vulnerability Fixer (CPU)")
+    gr.Markdown(
+        "Select your language, paste your code, and get a secured version of your code!"
+    )
+
+    with gr.Row():
+        with gr.Column():
+            lang_input = gr.Dropdown(
+                choices=["python", "java", "cpp"],
+                value="python",
+                label="Target Language",
+            )
+            code_input = gr.Textbox(
+                label="Vulnerable Code",
+                lines=15,
+                max_lines=30,
+                placeholder="Paste your vulnerable code here...",
+            )
+            submit_btn = gr.Button("Secure My Code ✨", variant="primary")
+
+        with gr.Column():
+            code_output = gr.Textbox(
+                label="Fixed Code",
+                lines=15,
+                max_lines=30,
+                interactive=False,
+            )
+
+    gr.Examples(
+        examples=EXAMPLES,
+        inputs=[lang_input, code_input],
     )
 
-    # Update syntax highlighting based on dropdown
-    def update_syntax(lang):
-        return gr.update(language=lang), gr.update(language=lang)
-    
-    lang_input.change(update_syntax, lang_input, [code_input, code_output])
     submit_btn.click(fix_code, [lang_input, code_input], code_output)
 
-demo.launch()
+demo.launch(ssr_mode=False)