import gradio as gr from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed from transformers import pipeline title = "CodeGen Generator" description = "This is a subspace to make code generation with [CodeGen](https://huggingface.co/Salesforce/codegen-16B-mono), it is used in a larger [space](https://huggingface.co/spaces/loubnabnl/Code-generation-models-v1) for model comparison. We use the 6.1B parameters model in this space." example = [ ["def print_hello_world():", 8, 0.6, 42], ["def get_file_size(filepath):", 24, 0.6, 42], ["def count_lines(filename):", 40, 0.6, 42], ["def count_words(filename):", 40, 0.6, 42]] tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-6B-mono") model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-6B-mono") def code_generation(gen_prompt, max_tokens, temperature=0.6, seed=42): set_seed(seed) pipe = pipeline("text-generation", model=model, tokenizer=tokenizer) generated_text = pipe(gen_prompt, do_sample=True, top_p=0.95, temperature=temperature, max_new_tokens=max_tokens)[0]['generated_text'] return generated_text iface = gr.Interface( fn=code_generation, inputs=[ gr.Textbox(lines=10, label="Input code"), gr.inputs.Slider( minimum=8, maximum=10000, step=1, default=8, label="Number of tokens to generate", ), gr.inputs.Slider( minimum=0, maximum=2.5, 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()