cleanup body generation function

app.py

@@ -1,12 +1,11 @@
 import gradio as gr
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import datasets
-import asyncio
 import numpy as np
 import torch
 from threading import Thread
 
-from utils.tree_utils import parse_functions, get_docstrings, grab_before_comments, line_chr2char, node_str_idx
+from utils.tree_utils import parse_functions, get_docstrings, grab_before_comments, line_chr2char, node_str_idx, replace_function
 from utils.html_utils import make_iframe, construct_embed
 PIPE = None
 
@@ -86,7 +85,7 @@ def grab_sample(sample_idx):
     # funcs = _parse_functions(sample_code)
     # func_identifiers = [f"{idx:2d}: {n.child_by_field_name('declarator').text.decode()}" for idx, n in enumerate(funcs)]
     # 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
+    return sample_pass, sample_code, sample_title, source_iframe#, gr.Dropdown.update(choices=func_identifiers) #, sample_title, sample_auhtor
 
 def _make_pipeline(model_cp = "Vipitis/santacoder-finetuned-Shadertoys-fine"): #bad default model for testing
     # if torch.cuda.is_available():
@@ -205,15 +204,14 @@ def alter_body(old_code, func_id, funcs_list: list, prompt="", temperature=0.2,
         old_code (str): The original code.
         func_node (Node): The node of the function to replace the body of.
         funcs_list (list): The list of all functions in the code.
-        prompt (str): The prompt(title) to use for generation.
-        temperature (float): The temperature to use for generation.
-        max_new_tokens (int): The maximum number of tokens to generate.
-        top_p (float): The top_p to use for generation.
-        repetition_penalty (float): The repetition_penalty to use for generation.
+        prompt (str): The prompt(title) to use for generation. defaults to "".
+        temperature (float): The temperature to use for generation. defaults to 0.2.
+        max_new_tokens (int): The maximum number of tokens to generate. defaults to 512.
+        top_p (float): The top_p to use for generation. defaults to 0.95.
+        repetition_penalty (float): The repetition_penalty to use for generation. defaults to 1.2.
         pipeline (Pipeline): The pipeline to use for generation.
     Returns:
         str: The altered code.
-        pipeline (Pipeline): The pipeline to update the state
     """
     if isinstance(func_id, str):
         print(f"{func_id=}")
@@ -226,22 +224,13 @@ def alter_body(old_code, func_id, funcs_list: list, prompt="", temperature=0.2,
     print(f"using for generation: {func_node=}")
     
     generation_kwargs = _combine_generation_kwargs(temperature, max_new_tokens, top_p, repetition_penalty)
-    
-    print(f"{pipeline=}") # check if default even loaded
-    if pipeline is None:
-        print("no pipeline found, loading default one")
-        pipeline = _make_pipeline("Vipitis/santacoder-finetuned-Shadertoys-fine")
 
     func_start_idx = line_chr2char(old_code, func_node.start_point[0], func_node.start_point[1])
     identifier_str = func_node.child_by_field_name("type").text.decode() + " " + func_node.child_by_field_name("declarator").text.decode() #func_start_idx:body_start_idx?
     body_node = func_node.child_by_field_name("body")
-    body_start_idx, body_end_idx = node_str_idx(body_node) #can cause index error, needs better testing!
-    # body_start_idx = line_chr2char(old_code, body_node.start_point[0], body_node.start_point[1])
-    # body_end_idx = line_chr2char(old_code, body_node.end_point[0], body_node.end_point[1])
-    print(f"{old_code[body_start_idx:body_end_idx]=}")
+    body_start_idx, body_end_idx = node_str_idx(body_node) 
     model_context = identifier_str # base case
-    # add any comments at the beginning of the function to the model_context
-    # second_child = func_node.child_by_field_name("body").children[1] #might error out?
+    
     docstring = get_docstrings(func_node) #might be empty?
     if docstring:
         model_context = model_context + "\n" + docstring
@@ -255,31 +244,23 @@ def alter_body(old_code, func_id, funcs_list: list, prompt="", temperature=0.2,
     generation = _run_generation(model_context, pipeline, generation_kwargs)
     for i in generation:
         print(f"{i=}")
-        yield model_context + i, pipeline #fix in between, do all the stuff in the end?
+        yield model_context + i #fix in between, do all the stuff in the end?
     generation = i[:] #seems to work
     print(f"{generation=}")
     ctx_with_generation = model_context + generation
-    print(f"{ctx_with_generation=}")
     try:
         #strip the body
         first_gened_func = parse_functions(ctx_with_generation)[0] # truncate generation to a single function?
     except IndexError:
         print("generation wasn't a full function.")
         altered_code = old_code[:func_start_idx] + model_context + generation + "//the generation didn't complete the function!\n" + old_code[body_end_idx:] #needs a newline to break out of the comment.
-        return altered_code, pipeline
-        # raise gr.Error(f"didn't generate a full function: {generation!r}]")
-    # print(f"{first_gened_func=}")
-    generated_body = first_gened_func.child_by_field_name("body").text.decode()
-    print(f"{generated_body=}")
-    altered_code = old_code[:func_start_idx] + identifier_str + generated_body + old_code[body_end_idx:]
-    print(f"{altered_code=}") #we get here successfully
-    yield altered_code, pipeline #yield once so it updates? -> works... gg but doesn't seem to do it for the dropdown
-    return altered_code, pipeline #never gets used by the code block? maybe I need to yield it first? but works in the ov_notebook
+        return altered_code
+    altered_code = replace_function(func_node, first_gened_func)
+    yield altered_code #yield once so it updates? -> works... gg but doesn't seem to do it for the dropdown
+    return altered_code #never gets used by the code block? maybe I need to yield it first? but works in the ov_notebook
 
 def list_dropdown(in_code): #only used for auto update, not on sample pick?
     funcs = parse_functions(in_code)
-    
-    # print(f"updating drop down to:{func_identifiers=}")
     func_identifiers = [f"{idx:2d}: {n.child_by_field_name('declarator').text.decode()}" for idx, n in enumerate(funcs)]
     # funcs = [n for n in funcs] #wrapped as set to avoid json issues?
     print(f"updating drop down to:{func_identifiers}")
@@ -349,17 +330,14 @@ if __name__ == "__main__": #works on huggingface?
             sample_code = gr.Code(new_shadertoy_code, label="Current Code (will update changes you generate)", language=None)
         bot_md = gr.Markdown(outro_text)
         sample_pass = gr.State(value={})
+        funcs = gr.State(value=[])
         pipe = gr.State(value=PIPE)
         pipe.value=_make_pipeline("Vipitis/santacoder-finetuned-Shadertoys-fine") # set a default like this? 
-        funcs = gr.State(value=[])
-        # funcs.value.append(list_dropdown(sample_code.value)[0]) #to circumvent the json issue?
-        # hist_state = gr.State(Value={})
-        # history_table = gr.JSON()
         
         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, 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(
+        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]).then(
             fn=list_dropdown, inputs=[sample_code], outputs=[funcs, func_dropdown]
         )
         sample_code.change(fn=list_dropdown, inputs=[sample_code], outputs=[funcs, func_dropdown]).then(