freealise/Code-Generation-with-Language-Specific-LoRa-Models

1import torch2import utils3import streamlit as st4import os5import subprocess6from datetime import datetime7 8 9def init_parameters():10    #Initialize the parameters11    # example_prompts_file_name = "example_prompts.json"12    example_codes_file_name = "example_codes.json"13    example_stop_tokens_file_name = "example_stop_tokens.json"14    # example_prompts = utils.read_json(example_prompts_file_name)15    example_codes = utils.read_json(example_codes_file_name)16    example_stop_tokens = utils.read_json(example_stop_tokens_file_name)17 18    java_example_prompts_file_name = "humaneval_java.jsonl"19    python_example_prompts_file_name = "humaneval_py.jsonl"20    ruby_example_prompts_file_name = "humaneval_rb.jsonl"21    rust_example_prompts_file_name = "humaneval_rs.jsonl"22    swift_example_prompts_file_name = "humaneval_swift.jsonl"23    java_example_prompts = utils.read_prompts(java_example_prompts_file_name)24    python_example_prompts = utils.read_prompts(python_example_prompts_file_name)25    ruby_example_prompts = utils.read_prompts(ruby_example_prompts_file_name)26    rust_example_prompts = utils.read_prompts(rust_example_prompts_file_name)27    swift_example_prompts = utils.read_prompts(swift_example_prompts_file_name)28    example_prompts = {29        "java": java_example_prompts,30        "python": python_example_prompts,31        "ruby": ruby_example_prompts,32        "rust": rust_example_prompts,33        "swift": swift_example_prompts34    }35    for key in example_prompts:36        if key not in example_stop_tokens:37            example_stop_tokens[key] = example_prompts[key]["prompt_stop_tokens"][0]38    return example_prompts, example_codes, example_stop_tokens39 40 41def get_programming_language():42    #Let the user choose the language between Python and Java43    lang = st.selectbox(44        "Choose the Programming Language in which you want to generate code",45        ("python", "java", "ruby", "rust", "swift")46    )47    return lang48 49 50def get_generation_stratgey(side_bar=True):51    #Let the user choose the generation strategy52    if side_bar:53        do_sample = st.sidebar.selectbox("do_sample: if set to True, this parameter enables decoding strategies such as multinomial sampling, beam-search multinomial sampling", (True, False))54        max_new_tokens = st.sidebar.number_input("max_new_tokens: The maximum number of tokens to generate. The higher this number, the longer the generation will take.", value=150)55        num_return_sequences = st.sidebar.number_input("num_return_sequences: The number of independently computed returned sequences for each element in the batch", value=1)56        temperature = st.sidebar.number_input("temperature: The value used to module the next token probabilities", value=0.2)57        top_p = st.sidebar.number_input("top_p: If set to float < 1, only the most probable tokens with probabilities that add up to top_p or higher are kept for generation", value=0.95)58    else:59        do_sample = st.selectbox("do_sample: if set to True, this parameter enables decoding strategies such as multinomial sampling, beam-search multinomial sampling", (True, False))60        max_new_tokens = st.number_input("max_new_tokens: The maximum number of tokens to generate. The higher this number, the longer the generation will take.", value=250)61        num_return_sequences = st.number_input("num_return_sequences: The number of independently computed returned sequences for each element in the batch", value=1)62        temperature = st.number_input("temperature: The value used to module the next token probabilities", value=0.2)63        top_p = st.number_input("top_p: If set to float < 1, only the most probable tokens with probabilities that add up to top_p or higher are kept for generation", value=0.95)64 65    gen_config_dict = {66        "do_sample": do_sample,67        "max_new_tokens": max_new_tokens,68        "num_return_sequences": num_return_sequences,69        "temperature": temperature,70        "top_p": top_p71    }72    gen = utils.initialize_generation_strategy_from_dict(gen_config_dict)73    return gen74 75 76def get_model_path(side_bar=True):77    #Let the user choose the Base Model  (wihout PEFT)78    base_model_paths = [79        'Salesforce/codegen-350M-mono',80        'ammarnasr/codegen-350M-mono-java',81        'ammarnasr/codegen-ruby-v7-run-1-checkpoint-100',82        'ammarnasr/codegen-350M-mono-rust',83        'ammarnasr/codegen-350M-mono-swift',84        85 86    ]87    base_model_paths_short = [88        'Baseline Mono',89        'Java LoRa',90        'Ruby LoRa',91        'Rust LoRa',92        'Swift LoRa',93    ]94 95    if side_bar:96        base_model_path = st.sidebar.selectbox("Choose the model for code compeletion", base_model_paths_short)97    else:98        base_model_path = st.selectbox("Choose the base model for code compeletion", base_model_paths_short)99 100    base_model_path = base_model_paths[base_model_paths_short.index(base_model_path)]101    return base_model_path102 103 104def get_device(side_bar=True):105    #Let the user choose the device106    opts = ["cpu"]107    if torch.cuda.is_available():108        opts.append("cuda")109    if side_bar:110        device = st.sidebar.selectbox("Choose the device",opts, index=len(opts)-1)111    else:112        device = st.selectbox("Choose the device",opts, index=len(opts)-1)113    return device114 115 116def code_generation_word_by_word(model, tokenizer, prompt, genration_stratgey, device, lang, STOP_TOKENS, tokens_per_iteration=1):117    """118    Generate code word by word and show the generated code in real time119    Args:120        model (torch.nn.Module): The model to use for code generation121        tokenizer (transformers.PreTrainedTokenizer): The tokenizer to use for tokenization122        prompt (str): The prompt to start the generation with123        genration_stratgey (transformers.GenerationStrategy): The generation strategy to use for generation124        device (str): The device to use for generation125        tokens_per_iteration (int, optional): The number of tokens to generate in each iteration. Defaults to 1.126    Returns:127        str: The generated code along with the prompt128    """129 130    # Intialize the parameters for real time code generation131    intial_prompt = prompt132    intial_prompt_len = len(intial_prompt)133    num_tokens_to_generate = genration_stratgey.max_new_tokens134    generated_tokens = 0135    genration_stratgey.max_new_tokens = tokens_per_iteration136    137    with st.empty(): # Set to empty to rewrite newly generated tokens inplace138        with torch.no_grad(): # Disable gradient calculation to reduce memory consumption139            while generated_tokens < num_tokens_to_generate: # Loop until the number of generated tokens is equal to the number of tokens to generate140                141                # For the first iteration, the inputs are the prompt, otherwise the inputs are the outputs of the previous iteration142                if generated_tokens == 0:143                    inputs = tokenizer(prompt, return_tensors="pt").to(device)144                    outputs = model.generate(input_ids=inputs.input_ids, attention_mask=inputs.attention_mask, generation_config=genration_stratgey)145                else:146                    outputs = model.generate(input_ids = outputs, generation_config=genration_stratgey)147 148                # Decode the generated tokens149                decoded_outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)150 151                # Add the decoded tokens to the prompt and show the prompt152                prompt += decoded_outputs[0][len(prompt):]153                st.code(prompt, language=lang)154                155                # Stop the generation if the generated tokens contain a stop token156                generated_text = prompt[intial_prompt_len:]157                generated_text_stopped = utils.stop_at_stop_token(generated_text, STOP_TOKENS)158                if generated_text_stopped != generated_text:159                    st.success("Code generated successfully")160                    prompt = intial_prompt + generated_text_stopped161                    break162                163                # Update the number of generated tokens164                generated_tokens += tokens_per_iteration165    return prompt166 167 168def load_model(model_path, device):169    #Load the model170    model_path_lower_case = model_path.lower()171    is_peft = False172    if "peft" in model_path_lower_case:173        is_peft = True174    if "lora" in model_path_lower_case:175        is_peft = True176    elif "ammar" in model_path_lower_case and "full" not in model_path_lower_case:177        is_peft = True178    if is_peft:179        model = utils.initialize_peft_model_from_huffingface(model_path)180    else:181        model = utils.initialize_causual_model_from_huffingface(model_path)182    model = model.to(device)183    return model184 185 186def write_current_solution_to_json(promt_and_code, example_prompts, rand_int, lang, genration_stratgey, edit_prompt=None):187    #Write the current solution to the json file188    prompt = example_prompts['prompt_text'][rand_int]189    if edit_prompt:190        code = promt_and_code[len(edit_prompt):]191    else:192        code = promt_and_code[len(prompt):]193    temp = genration_stratgey.temperature194    top_p = genration_stratgey.top_p195    max_new_tokens = genration_stratgey.max_new_tokens196    solution_dict = {197        "prompt": prompt,198        "tests": example_prompts['prompt_test'][rand_int],199        "stop_tokens": example_prompts['prompt_stop_tokens'][rand_int],200        "completions": [code],201        "temperature": temp,202        "top_p": top_p,203        "max_new_tokens": max_new_tokens,204        "language": lang,205    }206    current_soution_dir = "current_solution"207    if not os.path.exists(current_soution_dir):208        os.makedirs(current_soution_dir)209    current_solution_file_name = os.path.join(current_soution_dir, "current_solution.json")210    utils.write_json(current_solution_file_name, solution_dict)211 212    archive_dir = "archive"213    if not os.path.exists(archive_dir):214        os.makedirs(archive_dir)215    archive_file_name = os.path.join(archive_dir, f"current_solution_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.json")216    utils.write_json(archive_file_name, solution_dict)217 218 219def evalute_solution():220    td = 'current_solution'221    results_file = os.path.join(td, 'current_solution.results.json')222 223    #delete results file if exists224    if os.path.exists(results_file):225        os.remove(results_file)226 227    eval_cmd = f"podman run --rm --network none -v ./{td}:/{td}:rw multipl-e-eval --dir /{td} --output-dir /{td} --recursive"228    subprocess.run(eval_cmd.split())229    results = utils.read_json(results_file)230    st.write(results['results'][0]['status'])231    return results232 233 234def main():235    # set_page_config()236    col1, col2 = st.columns([3, 4])237    with col1:238        example_prompts, example_codes, example_stop_tokens = init_parameters()239        lang = get_programming_language()240        # example_codes = example_codes[lang]241        example_prompts = example_prompts[lang]242        STOP_TOKENS = example_stop_tokens[lang]243        device = get_device()244        model_path = get_model_path(side_bar=False)245        genration_stratgey = get_generation_stratgey()246        prompts_texts = example_prompts['prompt_text']247        rand_int = st.number_input("Choose a problem for the benchmark to solve (code below)", min_value=0, max_value=len(prompts_texts), value=50)248        default_prompt = prompts_texts[rand_int]249        # prompt = st.text_area("Enter the prompt to solve", value=default_prompt, height=200)250        prompt = default_prompt251        prompt_test = example_prompts['prompt_test'][rand_int]252        # prompt = prompt + "\n\n" + prompt_test253        st.code(prompt, language=lang)254        #Add tick box to edit prompt255        # edit_prompt = st.checkbox("Edit prompt", value=False)256        # if edit_prompt:257        #     prompt = st.text_area("Enter the prompt to solve", value=default_prompt, height=200)258        #     st.code(prompt, language=lang)259        # #Add tick box to enable/disable word by word generation260        # word_by_word_generation = st.checkbox("Word by word generation", value=True)261        edit_prompt = False262        word_by_word_generation = True263        # st.subheader("Generated Code")264        click = st.button("Generate the code")265    266    with col2:267        if click:268            with st.spinner("Generating the code ..."):269                if word_by_word_generation: # If the device is cuda, use the word by word generation strategy270                    tokenizer = utils.initialize_tokenizer_from_huggingface('Salesforce/codegen-350M-mono')271                    tokenizer.pad_token = tokenizer.eos_token272                    genration_stratgey.pad_token_id = tokenizer.pad_token_id273                    model = load_model(model_path, device)274                    promt_and_code = code_generation_word_by_word(model, tokenizer, prompt, genration_stratgey, device, lang, STOP_TOKENS)      275                else: # If the device is cpu, use the full generation strategy276                    st.info("loading the tokenizer ...")277                    tokenizer = utils.initialize_tokenizer_from_huggingface('Salesforce/codegen-350M-mono')278                    tokenizer.pad_token = tokenizer.eos_token279                    genration_stratgey.pad_token_id = tokenizer.pad_token_id280                    st.info("loading the model ...")281                    model = load_model(model_path, device)282                    st.info("tokenizing the prompt ...")283                    inputs = tokenizer(prompt, return_tensors="pt").to(device)284                    st.info("generating the code ...")285                    outputs = model.generate(**inputs, generation_config=genration_stratgey) 286                    st.info("decoding the code ...")287                    outputs = outputs[:, len(inputs["input_ids"][0]) :]288                    decoded_outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)289                    decoded_outputs = [utils.stop_at_stop_token(decoded_output, STOP_TOKENS) for decoded_output in decoded_outputs]290                    promt_and_code = prompt + "\n" + decoded_outputs[0] 291                # st.info("showing the generated code ...")292                st.code(promt_and_code, language=lang)    293                # st.info("writing the current solution to json ...")294                # write_current_solution_to_json(promt_and_code, example_prompts, rand_int, lang, genration_stratgey, edit_prompt=prompt)295                # # st.info("evaluating the current solution ...")296                # results = evalute_solution()297                # st.write(results)298                # program = results['results'][0]['program']299                # st.code(program, language=lang)300 301