amasad/Replit-v1-CodeInstruct-3B

1import os2import gradio as gr3import torch4 5from transformers import AutoTokenizer, AutoModelForCausalLM6 7REPO = "teknium/Replit-v1-CodeInstruct-3B"8 9description = """# <h1 style="text-align: center; color: white;"><span style='color: #F26207;'> Code Generation by Instruction with Replit-v1-CodeInstruct-3B </h1>10<span style="color: white; text-align: center;"> This model is trained on a large amount of code and fine tuned on code-instruct datasets. You can type an instruction in the ### Instruction: section and received code generation.</span>"""11 12device = "cuda" if torch.cuda.is_available() else "cpu"13 14tokenizer = AutoTokenizer.from_pretrained(REPO, trust_remote_code=True)15model = AutoModelForCausalLM.from_pretrained(REPO, torch_dtype=torch.bfloat16, trust_remote_code=True)16model.to(device)17 18model.eval()19 20custom_css = """21.gradio-container {22    background-color: #0D1525; 23    color:white24}25#orange-button {26    background: #F26207 !important;27    color: white;28}29.cm-gutters{30    border: none !important;31}32"""33 34def post_processing(prompt, completion):35    return prompt + completion36 37def code_generation(prompt, max_new_tokens=128, temperature=0.2, top_p=0.9, eos_token_id=tokenizer.eos_token_id):38    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)39    generated_ids = model.generate(input_ids, max_new_tokens=max_new_tokens, do_sample=True, use_cache=True, temperature=temperature, top_p=top_p, eos_token_id=eos_token_id)40    completion = tokenizer.decode(generated_ids[0][input_ids.shape[-1]:], skip_special_tokens=True, clean_up_tokenization_spaces=False)41    return post_processing(prompt, completion)42 43demo = gr.Blocks(44    css=custom_css45)46 47with demo:48    gr.Markdown(value=description)49    with gr.Row():50        input_col , settings_col  = gr.Column(scale=6), gr.Column(scale=6), 51        with input_col:52            code = gr.Code(lines=28,label='Input', value="### Instruction:\n\n### Response:")53        with settings_col:54            with gr.Accordion("Generation Settings", open=True):55                max_new_tokens= gr.Slider(56                    minimum=8,57                    maximum=128,58                    step=1,59                    value=48,60                    label="Max Tokens",61                )62                temperature = gr.Slider(63                    minimum=0.1,64                    maximum=2.5,65                    step=0.1,66                    value=0.2,67                    label="Temperature",68                )69 70    with gr.Row():71        run = gr.Button(elem_id="orange-button", value="Generate Response")72 73    event = run.click(code_generation, [code, max_new_tokens, temperature], code, api_name="predict")74 75demo.queue(max_size=40).launch()