fixes formatting

commit-0 · wenting-zhao · Dec 7, 2024 · Dec 8, 2024 · Dec 8, 2024 · Dec 8, 2024
commit 06daf7b48a828cb2d5e6713ebc407f4cc73523ad
diff --git a/commit0/harness/utils.py b/commit0/harness/utils.py
@@ -239,7 +239,7 @@ def extract_code_blocks(text: str) -> List[str]:
                    from the text.
 
     """
-    pattern = r"```python\n(.*?)```"
+    pattern = r'(?s)```(?:python|py)?(.*?)```'
     matches = re.finditer(pattern, text, re.DOTALL)
     return [match.group(1).strip() for match in matches]
 

diff --git a/examples/star/inference.py b/examples/star/inference.py
@@ -1,4 +1,4 @@
-from typing import List
+from typing import Dict, List
 from datasets import Dataset
 from vllm import LLM, SamplingParams
 from examples.star.utils import generate_prompt, cleanup
@@ -28,14 +28,14 @@ def generate_predictions(
     sampling_params = SamplingParams(n=n, temperature=temperature, max_tokens=512)
     llm = LLM(model=model_name)
 
-    prompts: List[str] = []
+    prompts: List[List[Dict]] = []
     for example in dataset:
         prompt = example["prompt"]
         test = example["test"]
         prompt = generate_prompt(prompt, test)
-        prompts.append(prompt)
+        prompts.append([{"role": "user", "content": prompt}])
 
-    outputs = llm.generate(prompts, sampling_params)
+    outputs = llm.chat(prompts, sampling_params)
 
     results: List[List[str]] = []
     for output in outputs:

diff --git a/pyproject.toml b/pyproject.toml
@@ -36,14 +36,15 @@ packages = ["commit0", "agent"]
 agent = [
     "import-deps>=0.3.0",
     "PyMuPDF>=1.24.5",
-    "aider-chat",
+    #"aider-chat>=0.69.1",
 ]
 example = [
-    "vllm",
-    "transformers",
+    "vllm==0.6.5",
+    "transformers>=4.47.1",
     "setuptools",
     "accelerate",
     "wandb>=0.19.0",
+    "openai>=1.58.1",
 ]
 
 [project.urls]