app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed
from transformers import pipeline


title = "InCoder Generator"
description = "This is a subspace to make code generation with [InCoder](https://huggingface.co./facebook/incoder-6B), it is used in a larger [space](https://huggingface.co./spaces/loubnabnl/Code-generation-models-v1) for model comparison."
example = [
    ["def count_words(filename):", 40, 0.6, 42],
    ["def print_hello_world():", "Sample", 8, 0.6, 42],
    ["def get_file_size(filepath):", "Sample", 22, 0.6, 42]]
tokenizer = AutoTokenizer.from_pretrained("facebook/incoder-1B")
model = AutoModelForCausalLM.from_pretrained("facebook/incoder-1B", low_cpu_mem_usage=True)
    

MAX_LENGTH = 2048
BOS = "<|endoftext|>"
def generate(gen_prompt, max_tokens, temperature=0.6, seed=42):
    set_seed(seed)
    input_ids = tokenizer(gen_prompt, return_tensors="pt").input_ids
    current_length = input_ids.flatten().size(0)
    max_length = max_tokens + current_length
    if max_length > MAX_LENGTH:
        max_length = MAX_LENGTH
    output = model.generate(input_ids=input_ids, do_sample=True, top_p=0.95, temperature=temperature, max_length=max_length)
    generated_text = tokenizer.decode(output.flatten())
    if generated_text.startswith(BOS):
        generated_text = detok_hypo_str[len(BOS):]
    return generated_text

generation = generate(gen_prompt, length_limit=40, temperature=0.6)

iface = gr.Interface(
    fn=generate, 
    inputs=[
        gr.Textbox(lines=10, label="Input code"),
        gr.inputs.Slider(
            minimum=8,
            maximum=256,
            step=1,
            default=8,
            label="Number of tokens to generate",
        ),
        gr.inputs.Slider(
            minimum=0,
            maximum=2,
            step=0.1,
            default=0.6,
            label="Temperature",
        )
        gr.inputs.Slider(
            minimum=0,
            maximum=1000,
            step=1,
            default=42,
            label="Random seed to use for the generation"
        )
    ],
    outputs=gr.Textbox(label="Predicted code", lines=10),
    examples=example,
    layout="horizontal",
    theme="peach",
    description=description,
    title=title
)
iface.launch()