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-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)."
example = [
    ["def count_words(filename):", 40, 0.6, 42],
    ["def print_hello_world():", 8, 0.6, 42],
    ["def get_file_size(filepath):", 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|>"
EXTENSION = "<| file ext=.py |>\n"

def generate(gen_prompt, max_tokens, temperature=0.6, seed=42):
    set_seed(seed)
    gen_prompt = EXTENSION + gen_prompt
    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 = generated_text[len(BOS):]
    generated_text = generated_text[len(EXTENSION):]
    return generated_text

iface = gr.Interface(
    fn=generate, 
    inputs=[
            gr.Code(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.1,
                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.Code(label="Predicted code", lines=10),
    examples=example,
    layout="horizontal",
    theme="peach",
    description=description,
    title=title
)
iface.launch()