Initial app

README.md

@@ -13,7 +13,8 @@ Check out the configuration reference at https://huggingface.co/docs/hub/spaces#
 
 
 ### Tasks that can be done 
-- Code Generation (Natural Language Query | Code Completion)
+- Code Generation (Natural Language Query | Code Completion) -> https://huggingface.co/lvwerra/codeparrot-small
 - Code Summarization (Summarize in Natural Language)
 - Code Search (from NL)
 - Similarity Check (Duplicate Detection)
+- Vulnerability Detection

app.py

@@ -1,7 +1,35 @@
 import gradio as gr
+from src.utils import return_parse_tree, GenCode, Config
+import logging
 
-def greet(name):
-    return "Hello " + name + "!!"
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+
+
+logger.info("Loading classes into the memory....")
+code_gen_class = GenCode(Config.code_gen_idt, gen_kwargs=Config.code_gen_config)
+logger.info("Sucessfully loaded classes into the memory")
+
+def code_gen(code_snippet:str) -> str:
+    logger.info("Starting to generate code...")
+    generated_code = code_gen_class(code_snippet)
+    logger.info("Sucessfully generated code")
+    
+    return generated_code
+
+def code_edit(code_snippet:str = "def dataset_name():\n    return 'test'"):
+    return code_snippet + "\n" + return_parse_tree(code_snippet)
+
+demo = gr.Blocks(analytics_enabled=True)
+with demo:
+    gr.Markdown("""# Code Analysis \n ## Code Editor""")
+    with gr.Row():
+        input_editor = gr.Textbox(label="Code Editor",lines=10,interactive=True,placeholder="Use this as the main code editor ")
+        button_gen = gr.Button(value="Complete Code")
+
+    button_gen.click(code_gen,input_editor,input_editor)
+    # iface = gr.Interface(title="Code Editor",fn=code_edit, inputs="text", outputs="text",examples=["def dataset_name():\n    return 'test'"])
+    # iface.launch()
+
+demo.launch(debug=True, server_port=8080)

flagged/log.csv

@@ -0,0 +1,3 @@
+'code_snippet','output','flag','username','timestamp'
+'','','','','2022-05-27 13:34:35.072204'
+'','','','','2022-05-27 13:34:35.721740'

src/test_utils.py

@@ -0,0 +1,12 @@
+from src.utils import GenCode, Config
+
+code_gen_class = GenCode(Config.code_gen_idt, gen_kwargs={})
+
+def test_gen_code():
+    code_snippet = "def dataset_name():\n    return 'test'"
+    generated_code = code_gen_class(code_snippet)
+    print(generated_code)
+
+
+if __name__ == "__main__":
+    test_gen_code()

src/utils.py

@@ -0,0 +1,26 @@
+import ast
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from dataclasses import dataclass
+
+# Configs
+class Config:
+    code_gen_idt : str = "lvwerra/codeparrot-small"
+    code_gen_config : dict = {
+        "max_length" : 256
+    }
+
+class GenCode:
+    def __init__(self,model_idt:str,gen_kwargs:dict) -> None:
+        self.generation_pipe = pipeline("text-generation", model=model_idt)
+        self.gen_kwargs = gen_kwargs
+    def __call__(self, input : str):
+        return self.generation_pipe(input, **self.gen_kwargs)[0]["generated_text"]
+        
+
+
+def return_lexer_map(code_snippet:str):
+    return None
+
+
+def return_parse_tree(code_snippet:str):
+    return ast.dump(ast.parse(code_snippet))