improve infilling script

ASSERT-KTH · Nov 19, 2024 · bafcd59 · bafcd59
1 parent 4f918ca
commit bafcd59
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/elleelleaime/generate/strategies/models/huggingface/codellama/codellama_infilling.py b/elleelleaime/generate/strategies/models/huggingface/codellama/codellama_infilling.py
@@ -16,9 +16,9 @@ class GenerateSettings:
     do_sample: bool = False
     temperature: float = 1.0
     num_beams: int = 1
-    early_stopping: bool = False
     num_return_sequences: int = 10
     max_length: int = 4096
+    early_stopping: bool = False
 
 
 class CodeLLaMAInfilling(PatchGenerationStrategy):
@@ -119,6 +119,7 @@ def __generate_patch(self, prompt: str) -> Optional[List[str]]:
                 early_stopping=self.generate_settings.early_stopping,
                 do_sample=self.generate_settings.do_sample,
                 temperature=self.generate_settings.temperature,
+                use_cache=True,
             )
 
         fillings_ids = generated_ids[:, input_len:]

diff --git a/generate_patches.py b/generate_patches.py
@@ -83,7 +83,7 @@ def entry_point(
             kwargs[key] = Path(kwargs[key]).name
 
     kwargs_str = "_".join([f"{k}={v}" for k, v in kwargs.items()])
-    kwargs_str = kwargs_str.replace("/", ":")
+    kwargs_str = kwargs_str.replace("/", "-")
     write_jsonl(
         os.path.join(
             dir_path,