fixed dropdown, imporve function generation

app.py

@@ -142,11 +142,13 @@ If I find an efficient way, the shaders might run in real time and be interactiv
 
 ## TODO:
  - [x] use embedded Shadertoy for reference/attribution (done, but some errors)
- - [] working render implementation on CPU only space (use the browser for WebGPU?, maybe via an iFrame too?)
- - [~] generate variations of return statements [ShaderEval task1](https://huggingface.co/spaces/Vipitis/ShaderEval) (missing all of the generation parameters)
+ - [] working render implementation on CPU only space (use the browser for WebGPU?, maybe via an iFrame too?) freespace uses lavapipe which works on really simple stuff.
+ - [~] generate variations of return statements [ShaderEval task1](https://huggingface.co/spaces/Vipitis/ShaderEval) (needs to be reworked using the other parts)
  - [] generation history stating which function and orig/generated returns. (use State ??). do it as comments in the code?
- - [x?] generate whole functions (only works once)
+ - [x] generate whole functions (seems to work quite well)
+ - [] display errros/issues to the user
  - [] generate whole shaders (via prompts?)
+ - [] accordion with generation parameters (as pipeline_kwargs?)
 """
 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.
@@ -184,11 +186,11 @@ def grab_sample(sample_idx):
     sample_auhtor = sample_pass["author"]
     source_iframe = construct_embed(sample_source)
     print(f"{source_iframe=}")
-    sample_funcs = _parse_functions(sample_code)
-    funcs = _parse_functions(sample_code)
-    func_identifiers = [(idx,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, source_iframe, funcs, gr.Dropdown.update(choices=func_identifiers) #, sample_title, sample_auhtor
+    # sample_funcs = _parse_functions(sample_code)
+    # 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, source_iframe, funcs#, gr.Dropdown.update(choices=func_identifiers) #, sample_title, sample_auhtor
 
 
 def _parse_functions(in_code):
@@ -280,7 +282,7 @@ def alter_body(old_code, func_id, funcs_list, pipeline=PIPE):
         str: The altered code.
     """
     print(f"{func_id=}")
-    func_id = int(func_id.split(",")[0]) #undo their string casting?
+    func_id = int(func_id.split(":")[0].strip()) #undo their string casting?
     func_node = funcs_list[func_id]
     print(f"using for generation: {func_node=}")
     
@@ -291,15 +293,22 @@ def alter_body(old_code, func_id, funcs_list, pipeline=PIPE):
         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()
     body_node = func_node.child_by_field_name("body")
     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]=}")
-    model_context = old_code[:body_start_idx]
-    generation = pipeline(model_context, max_new_tokens=(body_end_idx - body_start_idx)*2, return_full_text=False)[0]["generated_text"]
+    model_context = identifier_str # just this
+    num_new_tokens = (body_end_idx - body_start_idx) + 10 #TODO: approximation, we do have early stopping? maybe also use a number instead?
+    
+    print(f"generating up to {num_new_tokens} after {model_context!r}")
+    generation = pipeline(model_context, max_new_tokens=num_new_tokens, return_full_text=False)[0]["generated_text"]
     print(f"{generation=}")
-    first_gened_func = _parse_functions(old_code[func_start_idx:body_start_idx] + generation)[0] # truncate generation to a single function?
+    id_with_generation = identifier_str + generation
+    print(f"{id_with_generation=}")
+    first_gened_func = _parse_functions(id_with_generation)[0] # truncate generation to a single function?
     # strip just the body.
+    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[:body_start_idx] + generated_body + old_code[body_end_idx:]
@@ -310,11 +319,11 @@ def add_history(func_id, orig_rtn, gened_rtn, history):
     history[func_id] =  (orig_rtn, gened_rtn)
     return history, history
 
-def list_dropdown(in_code):
+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 = [(idx,n.child_by_field_name("declarator").text.decode()) for idx, n in enumerate(funcs)]
+    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}")
     return funcs, gr.Dropdown.update(choices=func_identifiers)
@@ -347,12 +356,13 @@ with gr.Blocks() as site:
     # history_table = gr.JSON()
     
     model_cp.submit(fn=_make_pipeline, inputs=[model_cp], outputs=[pipe])
-    sample_idx.release(fn=grab_sample, inputs=[sample_idx], outputs=[sample_pass, sample_code, source_embed, funcs, func_dropdown])
+    sample_idx.release(fn=grab_sample, inputs=[sample_idx], outputs=[sample_pass, sample_code, source_embed])
     # sample_idx.release(fn=list_dropdown, inputs=[sample_code], outputs=[funcs, func_dropdown]) #use multiple event handles to call other functions! seems to not work really well. always messes up
     gen_return_button.click(fn=alter_return, inputs=[sample_code, time_slider, pipe], outputs=[sample_code])
     gen_func_button.click(fn=alter_body, inputs=[sample_code, func_dropdown, funcs, pipe], outputs=[sample_code])
     # run_button.click(fn=add_history, inputs=[time_slider, sample_pass, sample_code, hist_state], outputs=[history_table, hist_state])
     # sample_idx.release(fn=construct_embed, inputs=[sample_idx], outputs=[source_embed]) #twice to make have different outputs?
+    sample_code.change(fn=list_dropdown, inputs=[sample_code], outputs=[funcs, func_dropdown]) # to update this after generation, so spans aren't messed up
     time_slider.release(fn=lambda code, time: asyncio.run(get_image(code, time)), inputs=[sample_code, time_slider], outputs=rendered_frame)
     render_button.click(fn=lambda code: asyncio.run(get_image(code)), inputs=[sample_code], outputs=rendered_frame)
     # run_button.click(fn=print, inputs=[model_cp, sample_idx], outputs=output)