Feature/add formatter (#473)

* Add formatter * Modify submission formatter
emrgnt-cmplxty · Aug 9, 2023 · 2e4a289 · 2e4a289
1 parent e58452e
commit 2e4a289
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 10 deletions.
diff --git a/automata-embedding-data b/automata-embedding-data
diff --git a/poetry.lock b/poetry.lock
diff --git a/research/study_leetcode/leetcode_problems_loader.py b/research/study_leetcode/leetcode_problems_loader.py
@@ -1,4 +1,4 @@
-from typing import Tuple
+from typing import Any, Tuple
 
 import pandas as pd
 
@@ -22,7 +22,7 @@ def get_problem_context(self, idx: int) -> str:
         row = self.data.iloc[idx]
         return f"Title:{row['question_title']}\n\nDescription:\n{row['description']}\n\nNote, your final solution MUST conform to the snippet shown here - ```python\\n{row['python3_snippet']}```"
 
-    def get_problem_id_slug(self, idx: int) -> Tuple[int, str]:
+    def get_problem_id_slug(self, idx: int) -> Tuple[int, int, Any]:
         """Retrieve a problem by its index."""
         row = self.data.iloc[idx]
         return (

diff --git a/research/study_leetcode/leetcode_solutions_finder.py b/research/study_leetcode/leetcode_solutions_finder.py
@@ -93,8 +93,8 @@ def find_best_match_and_explanation(self, query: str) -> str:
         self.solutions_data["similarity"] = self.solutions_data[
             "embedding"
         ].apply(
-            lambda x: self.calculate_similarity(x, context_embedding)
-        )  # type: ignore
+            lambda x: self.calculate_similarity(x, context_embedding)  # type: ignore
+        )
 
         # Sort solutions by similarity
         solutions_data_sorted = self.solutions_data.sort_values(

diff --git a/research/study_leetcode/run_automata_problem_solver.py b/research/study_leetcode/run_automata_problem_solver.py
@@ -50,6 +50,7 @@
     LeetCodeSubmission,
     ProgrammingLanguage,
 )
+from leetcode_env.utils import PySubmissionFormatter
 
 
 def main():  # sourcery skip: docstrings-for-functions
@@ -90,7 +91,7 @@ def main():  # sourcery skip: docstrings-for-functions
     print(f"Loading problem data from {args.problems_data_path}")
     loader = LeetCodeLoader(args.problems_data_path)
     embedding_provider = OpenAIEmbeddingProvider()
-
+    formatter = PySubmissionFormatter
     print(f"Number of examples to run = {len(loader.data)}")
     success_count = 0
     results = {}
@@ -183,7 +184,7 @@ def main():  # sourcery skip: docstrings-for-functions
             )
             lang = ProgrammingLanguage.PYTHON3
             sub = LeetCodeSubmission(
-                code=code,
+                code=formatter.to_leetcode(code),
                 lang=lang,
                 question_id=backend_problem_id,
                 question_slug=problem_slug,