refactor tree_utils

app.py

@@ -6,6 +6,9 @@ import numpy as np
 import torch
 from threading import Thread
 
+from tree_utils import _parse_functions, _get_docstrings, _grab_before_comments, _line_chr2char
+PIPE = None
+
 def make_script(shader_code):
     # code copied and fixed(escaping single quotes to double quotes!!!) from https://webglfundamentals.org/webgl/webgl-shadertoy.html
     script = ("""
@@ -295,18 +298,6 @@ new_shadertoy_code = """void mainImage( out vec4 fragColor, in vec2 fragCoord )
     fragColor = vec4(col,1.0);
 }"""
 
-passes_dataset = datasets.load_dataset("Vipitis/Shadertoys")
-single_passes = passes_dataset.filter(lambda x: not x["has_inputs"] and x["num_passes"] == 1) #could also include shaders with no extra functions.
-# single_passes = single_passes.filter(lambda x: x["license"] not in "copyright") #to avoid any "do not display this" license?
-all_single_passes = datasets.concatenate_datasets([single_passes["train"], single_passes["test"]])
-num_samples = len(all_single_passes)    
-
-import tree_sitter
-from tree_sitter import Language, Parser
-Language.build_library("./build/my-languages.so", ['tree-sitter-glsl'])
-GLSL_LANGUAGE = Language('./build/my-languages.so', 'glsl')
-parser = Parser()
-parser.set_language(GLSL_LANGUAGE)
 
 def grab_sample(sample_idx):
     sample_pass = all_single_passes[sample_idx]
@@ -322,19 +313,6 @@ def grab_sample(sample_idx):
     # print(f"updating drop down to:{func_identifiers}")
     return sample_pass, sample_code, sample_title, source_iframe, funcs#, gr.Dropdown.update(choices=func_identifiers) #, sample_title, sample_auhtor
 
-
-def _parse_functions(in_code):
-    """
-    returns all functions in the code as their actual nodes.
-    includes any comment made directly after the function definition or diretly after #copilot trigger
-    """
-    tree = parser.parse(bytes(in_code, "utf8"))
-    funcs = [n for n in tree.root_node.children if n.type == "function_definition"]
-
-    return funcs
-
-PIPE = None
-
 def _make_pipeline(model_cp = "Vipitis/santacoder-finetuned-Shadertoys-fine"): #bad default model for testing
     # if torch.cuda.is_available():
     #     device = "cuda"
@@ -436,16 +414,6 @@ def alter_return(orig_code, func_idx, temperature, max_new_tokens, top_p, repeti
     
     return altered_code
 
-def _line_chr2char(text, line_idx, chr_idx):
-    """
-    returns the character index at the given line and character index.
-    """
-    lines = text.split("\n")
-    char_idx = 0
-    for i in range(line_idx):
-        char_idx += len(lines[i]) + 1
-    char_idx += chr_idx
-    return char_idx
 
 def _combine_generation_kwargs(temperature, max_new_tokens, top_p, repetition_penalty):
     gen_kwargs = {}
@@ -455,34 +423,6 @@ def _combine_generation_kwargs(temperature, max_new_tokens, top_p, repetition_pe
     gen_kwargs["repetition_penalty"] = repetition_penalty
     return gen_kwargs
 
-def _grab_before_comments(func_node):
-    """
-    returns the comments that happen just before a function node
-    """
-    precomment = ""
-    last_comment_line = 0
-    for node in func_node.parent.children: #could you optimize where to iterated from? directon?
-        if node.start_point[0] != last_comment_line + 1:
-            precomment = ""
-        if node.type == "comment":
-            precomment += node.text.decode() + "\n"
-            last_comment_line = node.start_point[0]
-        elif node == func_node:
-            return precomment
-    return precomment
-
-def _get_docstrings(func_node):
-    """
-    returns the docstring of a function node
-    """
-    docstring = ""
-    for node in func_node.child_by_field_name("body").children:
-        if node.type == "comment" or node.type == "{":
-            docstring += node.text.decode() + "\n"
-        else:
-            return docstring
-    return docstring
-
 def alter_body(old_code, func_id, funcs_list: list, prompt, temperature, max_new_tokens, top_p, repetition_penalty, pipeline=PIPE):
     """
     Replaces the body of a function with a generated one.
@@ -581,7 +521,7 @@ def construct_embed(source_url):
 with gr.Blocks() as site:
     top_md = gr.Markdown(intro_text)
     model_cp = gr.Textbox(value="Vipitis/santacoder-finetuned-Shadertoys-fine", label="Model Checkpoint (Enter to load!)", interactive=True)
-    sample_idx = gr.Slider(minimum=0, maximum=num_samples, value=3211, label="pick sample from dataset", step=1.0)
+    sample_idx = gr.Slider(minimum=0, maximum=10513, value=3211, label="pick sample from dataset", step=1.0)
     func_dropdown = gr.Dropdown(value=["0: edit the Code (or load a shader) to update this dropdown"], label="chose a function to modify") #breaks if I add a string in before that? #TODO: use type="index" to get int - always gives None?
     prompt_text = gr.Textbox(value="the title used by the model has generation hint", label="prompt text", info="leave blank to skip", interactive=True)
     with gr.Accordion("Advanced settings", open=False): # from: https://huggingface.co/spaces/bigcode/bigcode-playground/blob/main/app.py
@@ -644,7 +584,7 @@ with gr.Blocks() as site:
     
     model_cp.submit(fn=_make_pipeline, inputs=[model_cp], outputs=[pipe]) # how can we trigger this on load?
     sample_idx.release(fn=grab_sample, inputs=[sample_idx], outputs=[sample_pass, sample_code, prompt_text, source_embed]) #funcs here?
-    gen_return_button.click(fn=alter_return, inputs=[sample_code, func_dropdown, pipe], outputs=[sample_code])
+    gen_return_button.click(fn=alter_return, inputs=[sample_code, func_dropdown, temperature, max_new_tokens, top_p, repetition_penalty, pipe], outputs=[sample_code])
     gen_func_button.click(fn=alter_body, inputs=[sample_code, func_dropdown, funcs, prompt_text, temperature, max_new_tokens, top_p, repetition_penalty, pipe], outputs=[sample_code, pipe]).then(
         fn=list_dropdown, inputs=[sample_code], outputs=[funcs, func_dropdown]
     )
@@ -652,5 +592,11 @@ with gr.Blocks() as site:
         fn=make_iframe, inputs=[sample_code], outputs=[our_embed])
 
 if __name__ == "__main__": #works on huggingface?
+    passes_dataset = datasets.load_dataset("Vipitis/Shadertoys")
+    single_passes = passes_dataset.filter(lambda x: not x["has_inputs"] and x["num_passes"] == 1) #could also include shaders with no extra functions.
+    # single_passes = single_passes.filter(lambda x: x["license"] not in "copyright") #to avoid any "do not display this" license?
+    all_single_passes = datasets.concatenate_datasets([single_passes["train"], single_passes["test"]])
+    num_samples = len(all_single_passes)    
+    
     site.queue()
     site.launch()

tree_utils.py

@@ -0,0 +1,59 @@
+import tree_sitter
+from tree_sitter import Language, Parser
+
+Language.build_library("./build/my-languages.so", ['tree-sitter-glsl'])
+GLSL_LANGUAGE = Language('./build/my-languages.so', 'glsl')
+parser = Parser()
+parser.set_language(GLSL_LANGUAGE)
+
+
+def _parse_functions(in_code):
+    """
+    returns all functions in the code as their actual nodes.
+    includes any comment made directly after the function definition or diretly after #copilot trigger
+    """
+    tree = parser.parse(bytes(in_code, "utf8"))
+    funcs = [n for n in tree.root_node.children if n.type == "function_definition"]
+
+    return funcs
+
+
+def _get_docstrings(func_node):
+    """
+    returns the docstring of a function node
+    """
+    docstring = ""
+    for node in func_node.child_by_field_name("body").children:
+        if node.type == "comment" or node.type == "{":
+            docstring += node.text.decode() + "\n"
+        else:
+            return docstring
+    return docstring
+
+
+def _grab_before_comments(func_node):
+    """
+    returns the comments that happen just before a function node
+    """
+    precomment = ""
+    last_comment_line = 0
+    for node in func_node.parent.children: #could you optimize where to iterated from? directon?
+        if node.start_point[0] != last_comment_line + 1:
+            precomment = ""
+        if node.type == "comment":
+            precomment += node.text.decode() + "\n"
+            last_comment_line = node.start_point[0]
+        elif node == func_node:
+            return precomment
+    return precomment
+
+def _line_chr2char(text, line_idx, chr_idx):
+    """
+    returns the character index at the given line and character index.
+    """
+    lines = text.split("\n")
+    char_idx = 0
+    for i in range(line_idx):
+        char_idx += len(lines[i]) + 1
+    char_idx += chr_idx
+    return char_idx