Merge pull request #129 from MLDSAI/fix/prompt_truncation

fix prompt truncation; fix diff_mask
OpenAdaptAI · May 9, 2023 · 0ab8213 · 0ab8213
2 parents 4f22e89 + dd045ef
commit 0ab8213
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/puterbot/models.py b/puterbot/models.py
@@ -197,7 +197,7 @@ def diff(self):
     @property
     def diff_mask(self):
         if not self._diff_mask:
-            self._diff_mask = self._diff.convert("1")
+            self._diff_mask = self.diff.convert("1")
         return self._diff_mask
 
     @property

diff --git a/puterbot/strategies/llm_mixin.py b/puterbot/strategies/llm_mixin.py
@@ -44,7 +44,11 @@ def get_completion(
             logger.warning(
                 f"Truncating from {len(prompt)=} to {max_input_size=}"
             )
-            prompt = prompt[max_input_size:]
+            prompt = prompt[-max_input_size:]
+            logger.warning(
+                f"Truncated {len(prompt)=}"
+            )
+
         logger.debug(f"{prompt=} {max_tokens=}")
         input_tokens = self.tokenizer(prompt, return_tensors="pt")
         pad_token_id = self.tokenizer.eos_token_id