Create app.py

app.py

@@ -0,0 +1,150 @@
+import streamlit as st
+from langchain_community.llms import OpenAI
+import argparse
+from datasets import load_dataset
+import yaml
+from tqdm import tqdm
+import re
+
+
+
+def load_data(split="test"):
+    data = load_dataset("bigcode/humanevalpack")
+    print("=========== dataset statistics ===========")
+    print(len(data[split]))
+    print("==========================================")
+    return data[split]
+
+
+
+def split_function_header_and_docstring(s):
+    # pattern = re.compile(r'\"\"\"(.*?)\"\"\"', re.DOTALL)
+    pattern = re.compile(r"(\"\"\"(.*?)\"\"\"|\'\'\'(.*?)\'\'\')", re.DOTALL)
+    match = pattern.findall(s)
+    if match:
+        # docstring = match.group(-1)
+        docstring = match[-1][0]
+        code_without_docstring = s.replace(docstring, "").replace('"' * 6, "").strip()
+        docstring = docstring.replace('"', "")
+    else:
+        raise ValueError
+    return code_without_docstring, docstring
+
+
+def prepare_model_input(code_data):
+    prompt = """Provide feedback on the errors in the given code and suggest the
+    correct code to address the described problem.
+
+Problem Description:
+{description}
+
+Incorrect Code:
+{wrong_code}"""
+
+    description = code_data["prompt"]
+    function_header, docstring = split_function_header_and_docstring(description)
+    problem = docstring.split(">>>")[0]
+
+    wrong_code = function_header + code_data["buggy_solution"]
+    template_dict = {"function_header": function_header, "description": problem, "wrong_code": wrong_code}
+    model_input = prompt.format(**template_dict)
+    return model_input, problem, function_header
+
+
+def load_and_prepare_data():
+    dataset = load_data()
+    all_model_inputs = {} 
+    print("### load and prepare data")
+    for data in tqdm(dataset):
+        problem_id = data['task_id'] 
+        buggy_solution = data['buggy_solution']
+        model_input, problem, function_header = prepare_model_input(data)
+        new_model_input =f"Provide feedback on the errors in the given code and suggest the correct code to address the described problem.\nProblem Description:{problem}\nIncorrect Code:\n{buggy_solution}\nFeedback:"
+        # data["header"] = function_header
+        all_model_inputs[problem_id] = {
+                "model_input": new_model_input,
+                "header": function_header,
+                "problem_description": problem,
+                "data": data
+            }
+    return all_model_inputs
+
+ 
+
+
+dataset = load_dataset("bigcode/humanevalpack", split='test',  trust_remote_code=True)  # Ensuring consistent split usage
+
+problem_ids = [problem['task_id'] for problem in dataset]
+all_model_inputs = load_and_prepare_data()
+
+
+# Initialize with dummy ports for demonstration purposes here
+parser = argparse.ArgumentParser()
+parser.add_argument("--editor_port", type=str, default="6000")
+parser.add_argument("--critic_port", type=str, default="6001")
+
+# Assuming args are passed via command line interface
+args = parser.parse_args()
+
+# Initialize Langchain LLMs for our models (please replace 'your_api_key' with actual API keys)
+editor_model = OpenAI(model="Anonymous-COFFEE/COFFEEPOTS-editor", api_key="EMPTY", openai_api_base=f"https://editor.jp.ngrok.io/v1")
+# critic_model = OpenAI(model="Anonymous-COFFEE/COFFEEPOTS-critic", api_key="EMPTY", openai_api_base=f"http://localhost:{args.critic_port}/v1")
+
+critic_model = OpenAI(model="Anonymous-COFFEE/COFFEEPOTS-critic", api_key="EMPTY", openai_api_base=f"https://critic.jp.ngrok.io/v1")
+
+st.title("Demo for COFFEEPOTS")
+
+selected_task_id = st.selectbox("Select a problem ID:", problem_ids)
+
+# Retrieve selected problem details
+problem_details = dataset[problem_ids.index(selected_task_id)]
+
+st.write(f"**Selected Problem ID:** {problem_details['task_id']}")
+st.write(f"**Problem Description:**\n{all_model_inputs[selected_task_id]['problem_description']}")
+# Display buggy code with syntax highlighting
+st.code(problem_details['buggy_solution'], language='python')
+
+status_text = st.empty()
+code_output = st.code("", language="python")
+
+def generate_feedback():
+
+    return critic_model.stream(input=f"{all_model_inputs[selected_task_id]['model_input']}", logit_bias=None)
+    # feedback = output.generations[0][0].text
+    # return feedback
+# def generate_corrected_code():
+#     return  "```python"+editor_model.stream(input=f"Buggy Code:\n{problem_details['buggy_solution']}\nFeedback: {feedback}", logit_bias=None)
+def generate_corrected_code():
+    # Stream output from the editor model
+    yield "```python"
+    for text_chunk in editor_model.stream(input=f"[INST]Buggy Code:\n{problem_details['buggy_solution']}\nFeedback: {feedback}[/INST]", logit_bias=None):
+        yield  text_chunk  # Assuming each chunk is part of the final code
+    
+    yield "```"
+        # time.sleep(0.02)  # Simulate processing delay; Adjust timing as necessary
+
+
+if st.button("Generate Feedback and Corrected Code"):
+    # Example of generating feedback and corrected code (replace these with actual model calls)
+    with st.spinner("Generating feedback..."):
+        # Simulate API call to critic_model
+        print(f"model input for critic:")
+        print(all_model_inputs[selected_task_id]['model_input'])
+        # output = critic_model.generate(prompts=[f"{all_model_inputs[selected_task_id]['model_input']}"], logit_bias=None)
+        # feedback = output.generations[0][0].text
+        # print(feedback)
+        # feedback = "dummy feedback"
+        
+        # status_text.markdown(f"{feedback}")
+        feedback = status_text.write_stream(generate_feedback())
+        # status_text.code(f"{feedback}", language='python')
+    
+    with st.spinner("Generating corrected code..."):
+        # Simulate API call to editor_model
+        # output = editor_model.generate(prompts=[f"Buggy Code:\n{problem_details['buggy_solution']}\nFeedback: {feedback}"], logit_bias=None)
+        # corrected_code = output.generations[0][0].text
+        # print(corrected_code)
+        # corrected_code = "dummy code"
+        # st.write("**Corrected Code:**")
+        corrected_code = code_output.write_stream(generate_corrected_code())
+        # code_output.code(corrected_code, language='python')