Add func to clean up the text generated by the model and added link to wiki

app.py

@@ -4,6 +4,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 # model_name = "flax-community/gpt-neo-1.3B-apps-all"
 model_name = "flax-community/gpt-neo-125M-apps-all"
 
+
 @st.cache(allow_output_mutation=True, max_entries=1)
 def get_model():
     model = AutoModelForCausalLM.from_pretrained(model_name)
@@ -13,16 +14,29 @@ def get_model():
 
 
 def format_input(question, starter_code=""):
-    answer_type = "\nUse Call-Based format\n" if starter_code else \
-                  "\nUse Standard Input format\n"
+    answer_type = (
+        "\nUse Call-Based format\n" if starter_code else "\nUse Standard Input format\n"
+    )
     return f"\nQUESTION:\n{question}\n{starter_code}\n{answer_type}\nANSWER:\n"
 
 
-def generate_solution(model, tokenizer, question, starter_code="", temperature=1.0, num_beams=1):
+def clean_text(generation):
+    # clean up text has discussed in OpenAI's paper "Evaluating Large Language Models Trained on Code"
+    generation = generation.split("\ndef")[0]
+    generation = generation.split("\nclass")[0]
+    generation = generation.split("\n#")[0]
+    generation = generation.split("\nif")[0]
+
+    return generation
+
+
+def generate_solution(
+    model, tokenizer, question, starter_code="", temperature=1.0, num_beams=1
+):
     prompt = format_input(question, starter_code)
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids
     start = len(input_ids[0])
-    
+
     output = model.generate(
         input_ids,
         max_length=start + 150,
@@ -37,8 +51,10 @@ def generate_solution(model, tokenizer, question, starter_code="", temperature=1
         repetition_penalty=None,
         num_return_sequences=None,
     )
+    output_str = tokenizer.decode(output[0][start:], skip_special_tokens=True).strip()
+    output_str = clean_text(output_str)
 
-    return tokenizer.decode(output[0][start:], skip_special_tokens=True).strip()
+    return output_str
 
 
 _EXAMPLES = [
@@ -76,10 +92,10 @@ def greet(name, owner):
         0.8,
     ],
 ]
+
+
 def run():
-    st.set_page_config(
-        page_title="Code Clippy Problem Solver"
-    )
+    st.set_page_config(page_title="Code Clippy Problem Solver")
     # sidebar
     st.sidebar.title("Code Clippy")
     st.sidebar.image(
@@ -87,9 +103,10 @@ def run():
         caption="(c) awesome Aimee Trevett",
     )
     st.sidebar.markdown("[Github](https://github.com/ncoop57/gpt-code-clippy)")
-    
+    st.sidebar.markdown("[Report](https://github.com/ncoop57/gpt-code-clippy/wiki)")
+
     st.sidebar.markdown("### Controls:")
-    
+
     temperature = st.sidebar.slider(
         "Temperature",
         min_value=0.5,
@@ -113,17 +130,17 @@ def run():
         help="Text description of the coding problem to be solved",
     )
     starter_code = st.text_input(
-        "Started code: ",
-        value="def greet(name):",
-        help="Optional starter code"
+        "Started code: ", value="def greet(name):", help="Optional starter code"
     )
     submit_button = st.button("Solve")
 
     if submit_button:
         st.text("Solution:")
-        output = generate_solution(model, tokenizer, question, starter_code, temperature, num_beams)
+        output = generate_solution(
+            model, tokenizer, question, starter_code, temperature, num_beams
+        )
         st.code(output, language="python")
-    
 
-if __name__=="__main__":
+
+if __name__ == "__main__":
     run()