better backoff default + chat template in hf-local

Genaios · May 2, 2024 · caa4c30 · caa4c30
1 parent b6025d9
commit caa4c30
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 4 deletions.
diff --git a/text_machina/src/models/hf_local.py b/text_machina/src/models/hf_local.py
@@ -30,6 +30,14 @@ def generate_completion(self, prompt: str, generation_config: Dict) -> str:
         This method is not used, since generations are done
         with batches using `generate_completions`.
         """
+
+        if self.model_config.api_type == CompletionType.CHAT:
+            prompt = self.tokenizer.apply_chat_template(
+                [{"role": "user", "content": prompt}],
+                add_generation_prompt=True,
+                tokenize=False,
+            )
+
         tokenized = self.tokenizer(
             prompt, truncation=True, padding=True, return_tensors="pt"
         )

diff --git a/text_machina/src/models/hf_remote.py b/text_machina/src/models/hf_remote.py
@@ -27,9 +27,7 @@ def __init__(self, model_config: ModelConfig):
         retry_adapter = HTTPAdapter(
             max_retries=Retry(
                 total=getattr(self.model_config, "max_retries", 5),
-                backoff_factor=getattr(
-                    self.model_config, "backoff_factor", 2
-                ),
+                backoff_factor=getattr(self.model_config, "backoff_factor", 2),
                 status_forcelist=[
                     code for code in requests.status_codes._codes if code != 200
                 ],

diff --git a/text_machina/version.py b/text_machina/version.py
@@ -1,6 +1,6 @@
 _MAJOR = "0"
 _MINOR = "2"
-_REVISION = "10"
+_REVISION = "11"
 
 VERSION_SHORT = "{0}.{1}".format(_MAJOR, _MINOR)
 VERSION = "{0}.{1}.{2}".format(_MAJOR, _MINOR, _REVISION)