fix iframe, add generation

app.py

@@ -141,9 +141,10 @@ In the near future there will be some buttons and sliders to generate variations
 If I find an efficient way, the shaders might run in real time and be interactive.
 
 ## TODO:
- - [~] use embedded Shadertoy for reference/attribution (seems to not on change??)
+ - [x] use embedded Shadertoy for reference/attribution (done, but some errors)
  - [] working render implementation on CPU only space (use the browser for WebGPU?)
- - [] generate variations of return statements (ShaderEval task1)
+ - [~] generate variations of return statements [ShaderEval task1](https://huggingface.co/spaces/Vipitis/ShaderEval) (missing all of the generation parameters)
+ - [] generation history stating which function and orig/generated returns. (use State ??).
  - [] generate whole functions
  - [] generate whole shaders (via prompts?)
 """
@@ -170,30 +171,95 @@ def grab_sample(sample_idx):
     sample_title = sample_pass["title"]
     sample_auhtor = sample_pass["author"]
     source_iframe = construct_embed(sample_source)
+    print(f"{source_iframe=}")
     return sample_pass, sample_code, source_iframe #, sample_title, sample_auhtor
 
-def _make_pipeline(model_cp):
+PIPE = None #gloabl var in CAPS indicates constant? so why are we changing it?
+
+def _make_pipeline(model_cp = "gpt2"): #bad default model for testing
     tokenizer = AutoTokenizer.from_pretrained(model_cp, trust_remote_code=True)
     model = AutoModelForCausalLM.from_pretrained(model_cp, trust_remote_code=True)
-    return pipeline("text-generation", model=model, tokenizer=tokenizer, trust_remote_code=True)
+    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, trust_remote_code=True)
+    PIPE = pipe # set the global?
+    print(f"loaded model {model_cp} as a pipline")
+    return pipe
+
+def process_retn(retn):
+    return retn.split(";")[0].strip()
+
+def get_full_replacement(orig_code, retn_start_idx, retn_end_idx, prediction) -> str:
+    """
+    Batches the generated return statement into the code and returns the full altered code.
+    """
+    print(f"{orig_code[retn_start_idx:retn_end_idx]=}")
+    generated = process_retn(prediction)
+    print(f"{generated=}")
+    variation = orig_code[:retn_start_idx] + generated + orig_code[retn_end_idx:]
+    return variation
+
+def alter_return(orig_code, func_idx=0, pipeline=PIPE): #default pipeline can't be passed as gloabl?
+    """
+    Replaces the return statement of a function with a generated one.
+    Args:
+        orig_code (str): The original code.
+        func_idx (int): The index of the function to replace the return statement of.
+        pipeline (Pipeline): The pipeline to use for generation.
+    Returns:
+        str: The altered code.
+    """
+    if pipeline is None:
+        print("no pipeline found, loading default one")
+        pipeline = _make_pipeline()
+
+    retrns = []
+    retrn_start_idx = orig_code.find("return")
+    while retrn_start_idx != -1:
+        retrn_end_idx = orig_code.find(";", retrn_start_idx)
+        retrns.append((retrn_start_idx, retrn_end_idx))
+        retrn_start_idx = orig_code.find("return", retrn_end_idx)
+    num_returns = len(retrns)
+    if num_returns == 0:
+        print("no return statement found, returning original code")
+        return orig_code
+    func_idx = int(max(0, min(func_idx, num_returns - 1))) #clamp to valid range, cast to int as a bodge.
+    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?!
+    altered_code = get_full_replacement(orig_code, retrn_start_idx+7, retrn_end_idx, pipe_generation)
+    
+    return altered_code
+
+def add_history(func_id, orig_rtn, gened_rtn, history):
+    # is this a list? or a JSON dict?
+    history[func_id] =  (orig_rtn, gened_rtn)
+    return history, history
 
 def construct_embed(source_url):
-    return f'<iframe width="640" height="360" frameborder="0" src="{source_url}?gui=true&t=0&paused=true&muted=true" allowfullscreen></iframe>'
+    shader_id = source_url.split("/")[-1]
+    return f'<iframe width="640" height="360" frameborder="0" src="https://www.shadertoy.com/embed/{shader_id}?gui=true&t=0&paused=true&muted=true" allowfullscreen></iframe>'
 
 with gr.Blocks() as site:
     text_md = gr.Markdown(text)
-    model_cp = gr.Textbox(value="Vipitis/santacoder-finetuned-Shadertoys", label="Model Checkpoint", interactive=True)
-    sample_idx = gr.Slider(minimum=0, maximum=num_samples, value=5, label="pick sample from dataset", step=1.0)
-    # run_button = gr.Button(label="generate code")
-    render_button = gr.Button("render frame0",label="render frame")
-    time_slider = gr.Slider(minimum=0, maximum=10, value=0, label="time (update on release)", step=0.02)
+    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)
+    run_button = gr.Button("generate a alternate return statement for one function", label="generate code")
+    render_button = gr.Button("render frame0 (can carsh the sapce on invalid shadercode)",label="render frame")
+    time_slider = gr.Slider(minimum=0, maximum=10, value=0, label="time (update on release, also used to pick other functions as a bodge)", step=0.02)
     #output = gr.Textbox(label="Output")
     rendered_frame = gr.Image(shape=(512, 420), label=f"rendered frame preview")
     # info_md = gr.Markdown(value="code_source", label="source URL for this shader", interactive=False)
-    source_embed = gr.HTML('<iframe width="640" height="360" frameborder="0" src="https://www.shadertoy.com/embed/XtcSDf?gui=true&t=0&paused=true&muted=true" allowfullscreen></iframe>')
-    sample_code = gr.Code(label="Sample Code", language=None, readonly=True, lines=20)
+    source_embed = gr.HTML('<iframe width="640" height="360" frameborder="0" src="https://www.shadertoy.com/embed/WsBcWV?gui=true&t=0&paused=true&muted=true" allowfullscreen></iframe>', label="How this shader originally renders")
+    sample_code = gr.Code(label="Current Code (will update changes you generate)", language=None, readonly=True, lines=20)
     sample_pass = gr.State(value={})
+    pipe = gr.State(value=PIPE)
+    # hist_state = gr.State(Value={})
+    # 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])
+    run_button.click(fn=alter_return, inputs=[sample_code, time_slider, 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?
     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)