codeparrot/incoder-subspace

1import gradio as gr2from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed3from transformers import pipeline4 5 6title = "InCoder Generator"7description = "This is a subspace to make code generation with [InCoder-1B](https://huggingface.co/facebook/incoder-1B), it is used in a larger [space](https://huggingface.co/spaces/loubnabnl/Code-generation-models-v1) for model comparison. You can find the original demo for InCoder [here](https://huggingface.co/spaces/facebook/incoder-demo)."8example = [9    ["def count_words(filename):", 40, 0.6, 42],10    ["def print_hello_world():", 8, 0.6, 42],11    ["def get_file_size(filepath):", 22, 0.6, 42]]12tokenizer = AutoTokenizer.from_pretrained("facebook/incoder-1B")13model = AutoModelForCausalLM.from_pretrained("facebook/incoder-1B", low_cpu_mem_usage=True)14    15 16MAX_LENGTH = 204817BOS = "<|endoftext|>"18EXTENSION = "<| file ext=.py |>\n"19 20def generate(gen_prompt, max_tokens, temperature=0.6, seed=42):21    set_seed(seed)22    gen_prompt = EXTENSION + gen_prompt23    input_ids = tokenizer(gen_prompt, return_tensors="pt").input_ids24    current_length = input_ids.flatten().size(0)25    max_length = max_tokens + current_length26    if max_length > MAX_LENGTH:27        max_length = MAX_LENGTH28    output = model.generate(input_ids=input_ids, do_sample=True, top_p=0.95, temperature=temperature, max_length=max_length)29    generated_text = tokenizer.decode(output.flatten())30    if generated_text.startswith(BOS):31        generated_text = generated_text[len(BOS):]32    generated_text = generated_text[len(EXTENSION):]33    return generated_text34 35iface = gr.Interface(36    fn=generate, 37    inputs=[38            gr.Code(lines=10, label="Input code"),39            gr.inputs.Slider(40                minimum=8,41                maximum=256,42                step=1,43                default=8,44                label="Number of tokens to generate",45            ),46            gr.inputs.Slider(47                minimum=0.1,48                maximum=2,49                step=0.1,50                default=0.6,51                label="Temperature",52            ),53            gr.inputs.Slider(54                minimum=0,55                maximum=1000,56                step=1,57                default=42,58                label="Random seed to use for the generation"59            )60    ],61    outputs=gr.Code(label="Predicted code", lines=10),62    examples=example,63    layout="horizontal",64    theme="peach",65    description=description,66    title=title67)68iface.launch()