Spaces:
Runtime error
Runtime error
File size: 2,735 Bytes
8a3ef58 012c551 50c1955 012c551 0a4dd43 012c551 491ed03 012c551 8a3ef58 012c551 8a3ef58 012c551 8a3ef58 012c551 280eea5 012c551 491ed03 012c551 491ed03 012c551 491ed03 012c551 491ed03 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
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["do_sample"] = True
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=45.0) #IPEX takes a bit on first inference, to avoid an error with the empty queue timeout on the first time, we just wait longer.
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=""):
"""
Constructs the model context from a function node.
returns: model_context, start_byte
"""
model_context, start_byte = grab_before_comments(func_node)
model_context += full_func_head(func_node)
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, start_byte |