import gradio as gr from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed, pipeline title = "SantaCoder 🎅 Swift 🍏 Completion" description = "This is a subspace to make code generation with [SantaCoder fine-tuned on The Stack Swift](https://huggingface.co/mrm8488/santacoder-finetuned-the-stack-swift)" EXAMPLE_0 = "import SwiftUI\n\nstruct ContentView: View {\n var body: some View {" EXAMPLE_1 = "// Make a naviagtion list with the days of the week\nNavigationView {" CKPT = "mrm8488/santacoder-finetuned-the-stack-swift" examples = [[EXAMPLE_0, 9, 0.6, 42], [EXAMPLE_1, 114, 0.6, 42]] tokenizer = AutoTokenizer.from_pretrained(CKPT) model = AutoModelForCausalLM.from_pretrained(CKPT, trust_remote_code=True).to("cuda") def code_generation(gen_prompt, max_tokens, temperature=0.6, seed=42): set_seed(seed) pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0) 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=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=examples, layout="horizontal", theme="peach", description=description, title=title ) iface.launch()