from accelerate import Accelerator from transformers import TextIteratorStreamer from threading import Thread from .tree_utils import full_func_head, grab_before_comments def combine_generation_kwargs(temperature=2.0, max_new_tokens=512, top_p=0.95, repetition_penalty=1.2): """ Combines the generation kwargs into a single dict. """ 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 stream_generation(prompt:str, pipe, gen_kwargs:dict): accelerator = Accelerator() device = accelerator.device """ Text generation function Args: prompt (str): The context to start generation from. pipe (Pipeline): The pipeline to use for generation (we take the model and tokenizer form it) gen_kwargs (dict): The generation kwargs. Returns: str: The generated text. (it iterates over time) """ # Tokenize the model_context model_inputs = pipe.tokenizer(prompt, return_tensors="pt") model_inputs.to(device) model = pipe.model.to(device) #is this also required? # Start generation on a separate thread, so that we don't block the UI. The text is pulled from the streamer # in the main thread. Adds timeout to the streamer to handle exceptions in the generation thread. streamer = TextIteratorStreamer(pipe.tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15.0) generate_kwargs = dict(model_inputs, streamer=streamer, **gen_kwargs) t = Thread(target=pipe.model.generate, kwargs=generate_kwargs) t.start() # Pull the generated text from the streamer, and update the model output. model_output = "" for new_text in streamer: # print("step", end="") model_output += new_text yield model_output streamer.on_finalized_text("stream reached the end.") return model_output #is this ever reached? def construct_model_context(func_node, prompt="") -> str: """ Constructs the model context from a function node. """ model_context = grab_before_comments(func_node) + full_func_head(func_node) # (identifier + docstrings) if prompt != "": model_context = "//Title: " + prompt + "\n" + model_context #prepend user prompt/title model_context = "//Language: Shadertoy GLSL fragment shader\n" + model_context #prepend system prompt, language hint return model_context