|
import gradio as gr |
|
from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed |
|
from transformers import pipeline |
|
import os |
|
import torch |
|
|
|
description = """# <p style="text-align: center; color: white;"> π
<span style='color: #ff75b3;'>SantaCoder:</span> Code Generation </p> |
|
<span style='color: white;'>This is a demo to generate code with <a href="https://huggingface.co/bigcode/santacoder" style="color: #ff75b3;">SantaCoder</a>, |
|
a 1.1B parameter model for code generation in Python, Java & JavaScript. The model can also do infilling, just specify where you would like the model to complete code |
|
with the <span style='color: #ff75b3;'><FILL-HERE></span> token.</span>""" |
|
|
|
token = os.environ["HUB_TOKEN"] |
|
device="cuda:0" |
|
|
|
|
|
FIM_PREFIX = "<fim-prefix>" |
|
FIM_MIDDLE = "<fim-middle>" |
|
FIM_SUFFIX = "<fim-suffix>" |
|
FIM_PAD = "<fim-pad>" |
|
EOD = "<|endoftext|>" |
|
|
|
GENERATION_TITLE= "<p style='font-size: 16px; color: white;'>Generated code:</p>" |
|
|
|
tokenizer_fim = AutoTokenizer.from_pretrained("bigcode/santacoder", use_auth_token=token, padding_side="left") |
|
|
|
tokenizer_fim.add_special_tokens({ |
|
"additional_special_tokens": [EOD, FIM_PREFIX, FIM_MIDDLE, FIM_SUFFIX, FIM_PAD], |
|
"pad_token": EOD, |
|
}) |
|
|
|
tokenizer = AutoTokenizer.from_pretrained("bigcode/christmas-models", use_auth_token=token) |
|
model = AutoModelForCausalLM.from_pretrained("bigcode/christmas-models", trust_remote_code=True, use_auth_token=token).to(device) |
|
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device) |
|
|
|
def post_processing(prompt, completion): |
|
completion = "<span style='color: #ff75b3;'>" + completion + "</span>" |
|
prompt = "<span style='color: #727cd6;'>" + prompt + "</span>" |
|
code_html = f"<br><hr><br><pre style='font-size: 12px'><code>{prompt}{completion}</code></pre><br><hr>" |
|
return GENERATION_TITLE + code_html |
|
|
|
def post_processing_fim(prefix, middle, suffix): |
|
prefix = "<span style='color: #727cd6;'>" + prefix + "</span>" |
|
middle = "<span style='color: #ff75b3;'>" + middle + "</span>" |
|
suffix = "<span style='color: #727cd6;'>" + suffix + "</span>" |
|
code_html = f"<br><hr><br><pre style='font-size: 12px'><code>{prefix}{middle}{suffix}</code></pre><br><hr>" |
|
return GENERATION_TITLE + code_html |
|
|
|
def fim_generation(prompt, max_new_tokens, temperature): |
|
prefix = prompt.split("<FILL-HERE>")[0] |
|
suffix = prompt.split("<FILL-HERE>")[1] |
|
[middle] = infill((prefix, suffix), max_new_tokens, temperature) |
|
return post_processing_fim(prefix, middle, suffix) |
|
|
|
def extract_fim_part(s: str): |
|
|
|
start = s.find(FIM_MIDDLE) + len(FIM_MIDDLE) |
|
stop = s.find(EOD, start) or len(s) |
|
return s[start:stop] |
|
|
|
def infill(prefix_suffix_tuples, max_new_tokens, temperature): |
|
if type(prefix_suffix_tuples) == tuple: |
|
prefix_suffix_tuples = [prefix_suffix_tuples] |
|
|
|
prompts = [f"{FIM_PREFIX}{prefix}{FIM_SUFFIX}{suffix}{FIM_MIDDLE}" for prefix, suffix in prefix_suffix_tuples] |
|
|
|
inputs = tokenizer_fim(prompts, return_tensors="pt", padding=True, return_token_type_ids=False).to(device) |
|
with torch.no_grad(): |
|
outputs = model.generate( |
|
**inputs, |
|
do_sample=True, |
|
temperature=temperature, |
|
max_new_tokens=max_new_tokens, |
|
pad_token_id=tokenizer.pad_token_id |
|
) |
|
|
|
return [ |
|
extract_fim_part(tokenizer_fim.decode(tensor, skip_special_tokens=False)) for tensor in outputs |
|
] |
|
|
|
|
|
def code_generation(prompt, max_new_tokens, temperature=0.2, seed=42): |
|
|
|
|
|
if "<FILL-HERE>" in prompt: |
|
return fim_generation(prompt, max_new_tokens, temperature=0.2) |
|
else: |
|
completion = pipe(prompt, do_sample=True, top_p=0.95, temperature=temperature, max_new_tokens=max_new_tokens)[0]['generated_text'] |
|
completion = completion[len(prompt):] |
|
return post_processing(prompt, completion) |
|
|
|
|
|
demo = gr.Blocks( |
|
css=".gradio-container {background-color: #20233fff; color:white}" |
|
) |
|
with demo: |
|
with gr.Row(): |
|
_, colum_2, _ = gr.Column(scale=1), gr.Column(scale=6), gr.Column(scale=1) |
|
with colum_2: |
|
gr.Markdown(value=description) |
|
code = gr.Code(lines=5, language="python", label="Input code", value="def all_odd_elements(sequence):\n \"\"\"Returns every odd element of the sequence.\"\"\"") |
|
|
|
with gr.Accordion("Advanced settings", open=False): |
|
max_new_tokens= gr.Slider( |
|
minimum=8, |
|
maximum=1024, |
|
step=1, |
|
value=48, |
|
label="Number of tokens to generate", |
|
) |
|
temperature = gr.Slider( |
|
minimum=0.1, |
|
maximum=2.5, |
|
step=0.1, |
|
value=0.2, |
|
label="Temperature", |
|
) |
|
seed = gr.Slider( |
|
minimum=0, |
|
maximum=1000, |
|
step=1, |
|
label="Random seed to use for the generation" |
|
) |
|
run = gr.Button() |
|
output = gr.HTML(label="Generated code") |
|
|
|
event = run.click(code_generation, [code, max_new_tokens, temperature, seed], output, api_name="predict") |
|
gr.HTML(label="Contact", value="<img src='https://huggingface.co/datasets/bigcode/admin/resolve/main/bigcode_contact.png' alt='contact' style='display: block; margin: auto; max-width: 800px;'>") |
|
|
|
demo.launch() |