reshinthadith/code-brush

1import streamlit as st2from transformers import AutoTokenizer, AutoModelForCausalLM3import difflib4import re5from utils import verify_diff, apply_diff_from_output6 7commit_message_per_brush = {8    "Annotate Type": "annotate type to the variables.",9    "Reformat" : "Reformat the code using pep8",10    "Add Docstrings" : "Add docstrings to all the functions",11    "Add Comments" : "Add inline comments to all the functions",12}13 14 15def load_model_and_tokenizer(model_name:str="CarperAI/diff-codegen-350M-v2"):16    tokenizer = AutoTokenizer.from_pretrained(model_name)17    model = AutoModelForCausalLM.from_pretrained(model_name)18    return tokenizer, model19 20def make_prompt(code:str,task):21    filename = "input.py"22    prompt = f"<NME>main.py<BEF>{code}<MSG>{commit_message_per_brush[task]}."23    return prompt24 25 26def generate_diff(code:str):27    input_ids = tokenizer.encode(code, return_tensors='pt')28    outputs = model.generate(input_ids, max_length=64,temperature=0.8,top_p=0.85)29    return tokenizer.decode(outputs[0], skip_special_tokens=True)30 31 32def postprocess_output(generated_output:str):33    return verify_diff(generated_output)34 35st.title("Code Brush")36st.write("A tool to brush up your code")37 38 39 40tokenizer,model = load_model_and_tokenizer()41with st.form("my_form"):42    text = st.text_area("Enter your code here", height=150, value="def greet(input_name):\n    return f'Hello, {input_name}'" )43    brush_type = st.selectbox("Brush Type", ["Annotate Type", "Reformat", "Add Docstrings", "Add Comments"])44    submit_button = st.form_submit_button("Submit")45    if submit_button:46        st.write("## Diff:")47        generate_diff = generate_diff(make_prompt(text,brush_type))48        after_file = apply_diff_from_output(generate_diff)49        generate_diff_processed = postprocess_output(generate_diff)50        st.write(after_file)51        st.write(generate_diff_processed)52        #st.text_area(generate_diff_processed)53        #st.text_area(generate_diff, height=150, value=generate_diff)54 55    56 57 58 59