init commit

app.py

@@ -0,0 +1,51 @@
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import difflib
+import re
+
+commit_message_per_brush = {
+    "Annotate Type": "annotate type to the variables.",
+    "Reformat" : "Reformat the code using pep8",
+    "Add Docstrings" : "Add docstrings to all the functions",
+    "Add Comments" : "Add inline comments to all the functions",
+}
+
+
+def load_model_and_tokenizer(model_name:str="CarperAI/diff-codegen-2B-v2"):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    return tokenizer, model
+
+def make_prompt(code:str,task):
+    filename = "input.py"
+    prompt = f"<NME>main.py<BEF>{code}<MSG>{commit_message_per_brush[task]}."
+    return prompt
+
+
+def generate_diff(code:str):
+    input_ids = tokenizer.encode(code, return_tensors='pt')
+    outputs = model.generate(input_ids, max_length=64,temperature=0.8,top_p=0.85)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+
+def postprocess_output(generated_output:str):
+    pass
+
+st.title("Code Brush")
+st.write("A tool to brush up your code")
+
+
+
+tokenizer,model = load_model_and_tokenizer()
+with st.form("my_form"):
+    text = st.text_area("Enter your code here", height=150, value="def greet(input_name):\n    return f'Hello, {input_name}'" )
+    brush_type = st.selectbox("Brush Type", ["Annotate Type", "Reformat", "Add Docstrings", "Add Comments"])
+    submit_button = st.form_submit_button("Submit")
+    if submit_button:
+        st.write("## Diff:")
+        st.text_area(generate_diff(make_prompt(text,brush_type)))
+
+    
+
+
+

requirements.txt

@@ -0,0 +1,2 @@
+streamlit
+transformers