generation parameters added

app.py

@@ -269,7 +269,7 @@ outro_text ="""
  - [] generation history stating which function and orig/generated returns. (use State ??). do it as comments in the code?
  - [~] display errros/issues to the user (raise gr.Error could be one idea, but highlighting in the code would be awesome) currently adds a comment to the code.
  - [] generate whole shaders (via prompts guidance, recursive from errors)
- - [] accordion with generation parameters (as pipeline_kwargs?) look up starcoder playround and take "inspiration" from there
+ - [x] accordion with generation parameters (as pipeline_kwargs?) look up starcoder playround and take "inspiration" from there (implemented for both buttons, untested)
  - [] support FIM task for better model context
  - [~] include some context for prompt (title, comments before a functions) - now works with the first comment inside a function body (has to be first)
  - [] gradio examples
@@ -356,7 +356,7 @@ def get_full_replacement(orig_code, retn_start_idx, retn_end_idx, prediction) ->
     variation = orig_code[:retn_start_idx] + generated + orig_code[retn_end_idx:]
     return variation
 
-def alter_return(orig_code, func_idx, pipeline=PIPE): #default pipeline can't be passed as gloabl?
+def alter_return(orig_code, func_idx, temperature, max_new_tokens, top_p, repetition_penalty, pipeline=PIPE): #default pipeline can't be passed as gloabl?
     """
     Replaces the return statement of a function with a generated one.
     Args:
@@ -378,6 +378,8 @@ def alter_return(orig_code, func_idx, pipeline=PIPE): #default pipeline can't be
     else:
         raise gr.Error(f"func_idx must be int or str, not {type(func_idx)}")
 
+    generation_kwargs = _combine_generation_kwargs(temperature, max_new_tokens, top_p, repetition_penalty)
+
     retrns = []
     retrn_start_idx = orig_code.find("return")
     while retrn_start_idx != -1:
@@ -392,8 +394,7 @@ def alter_return(orig_code, func_idx, pipeline=PIPE): #default pipeline can't be
     retrn_start_idx, retrn_end_idx = retrns[func_idx]
     model_context = orig_code[:retrn_start_idx] #TODO: maximal context?
     model_inp = model_context + "return"
-    new_toks = (retrn_end_idx - retrn_start_idx) * 2 #TODO: approximation, we do have early stopping? maybe also use a number instead?
-    pipe_generation = pipeline(model_inp, max_new_tokens=new_toks, return_full_text=False)[0]["generated_text"] #pipeline kwargs are missing?!
+    pipe_generation = pipeline(model_inp, return_full_text=False, **generation_kwargs)[0]["generated_text"] #pipeline kwargs are missing?!
     altered_code = get_full_replacement(orig_code, retrn_start_idx+7, retrn_end_idx, pipe_generation)
     
     return altered_code
@@ -409,8 +410,15 @@ def _line_chr2char(text, line_idx, chr_idx):
     char_idx += chr_idx
     return char_idx
 
+def _combine_generation_kwargs(temperature, max_new_tokens, top_p, repetition_penalty):
+    gen_kwargs = {}
+    gen_kwargs["temperature"] = temperature
+    gen_kwargs["max_new_tokens"] = max_new_tokens
+    gen_kwargs["top_p"] = top_p
+    gen_kwargs["repetition_penalty"] = repetition_penalty
+    return gen_kwargs
 
-def alter_body(old_code, func_id, funcs_list: list, pipeline=PIPE):
+def alter_body(old_code, func_id, funcs_list: list, temperature, max_new_tokens, top_p, repetition_penalty, pipeline=PIPE):
     """
     Replaces the body of a function with a generated one.
     Args:
@@ -430,6 +438,7 @@ def alter_body(old_code, func_id, funcs_list: list, pipeline=PIPE):
     func_node = funcs_list[func_id]
     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:
@@ -449,10 +458,7 @@ def alter_body(old_code, func_id, funcs_list: list, pipeline=PIPE):
         # print(second_child.text.decode())
         model_context += " { \n  " + second_child.text.decode()
         print(f"{model_context=}")
-    num_new_tokens = max(160,(body_end_idx - body_start_idx) + 10) #TODO: approximation, we do have early stopping? maybe also use a number instead? HARD MAX for performance limits.
-    
-    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"]
+    generation = pipeline(model_context, return_full_text=False, **generation_kwargs)[0]["generated_text"]
     print(f"{generation=}")
     ctx_with_generation = model_context + generation
     print(f"{ctx_with_generation=}")
@@ -493,10 +499,50 @@ with gr.Blocks() as site:
     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)
     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?
+    with gr.Accordion("Advanced settings", open=False): # from: https://huggingface.co/spaces/bigcode/bigcode-playground/blob/main/app.py
+        with gr.Row():
+            column_1, column_2 = gr.Column(), gr.Column()
+            with column_1:
+                temperature = gr.Slider(
+                    label="Temperature",
+                    value=0.0, #start out at 0 to do greedy? or will there be an error?
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.05,
+                    interactive=True,
+                    info="Higher values produce more diverse outputs",
+                )
+                max_new_tokens = gr.Slider(
+                    label="Max new tokens",
+                    value=160,
+                    minimum=0,
+                    maximum=2048, #this could be inferred from the model?
+                    step=32,
+                    interactive=True,
+                    info="The maximum numbers of new tokens",
+                )
+            with column_2:
+                top_p = gr.Slider(
+                    label="Top-p (nucleus sampling)",
+                    value=0.30,
+                    minimum=0.0,
+                    maximum=1,
+                    step=0.05,
+                    interactive=True,
+                    info="Higher values sample more low-probability tokens",
+                )
+                repetition_penalty = gr.Slider(
+                    label="Repetition penalty",
+                    value=1.2,
+                    minimum=1.0,
+                    maximum=2.0,
+                    step=0.05,
+                    interactive=True,
+                    info="Penalize repeated tokens",
+                )
     with gr.Row():
         gen_return_button = gr.Button("generate a alternate return statement", label="generate return")
         gen_func_button = gr.Button("generate an alternate function body", label="generate function")
-        # update_funcs_button = gr.Button("update functions", label="update functions")
     with gr.Row():
         with gr.Column():
             source_embed = gr.HTML('<iframe width="640" height="360" frameborder="0" src="" allowfullscreen></iframe>', label="How this shader originally renders")
@@ -514,7 +560,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, source_embed]) 
     gen_return_button.click(fn=alter_return, inputs=[sample_code, func_dropdown, pipe], outputs=[sample_code])
-    gen_func_button.click(fn=alter_body, inputs=[sample_code, func_dropdown, funcs, pipe], outputs=[sample_code, pipe])
+    gen_func_button.click(fn=alter_body, inputs=[sample_code, func_dropdown, funcs, temperature, max_new_tokens, top_p, repetition_penalty, pipe], outputs=[sample_code, pipe])
     sample_code.change(fn=list_dropdown, inputs=[sample_code], outputs=[funcs, func_dropdown]) # to update this after generation, so spans aren't messed up
     sample_code.change(fn=make_iframe, inputs=[sample_code], outputs=[our_embed]) #twice could cause issues, find better ways.
 site.launch()