Update hub_model.py

promptslab · Jul 29, 2023 · bd2847c · bd2847c
1 parent 1cfa1a6
commit bd2847c
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/promptify/models/text2text/api/hub_model.py b/promptify/models/text2text/api/hub_model.py
@@ -32,7 +32,7 @@ def __init__(
         do_sample: bool = False,
         api_wait= 60,
         api_retry= 6,
-        max_completion_length: int = 20,
+        json_depth_limit: int = 20,
     ):
 
 
@@ -50,7 +50,7 @@ def __init__(
         self.max_time = max_time
         self.num_return_sequences = num_return_sequences
         self.do_sample = do_sample
-        self.max_completion_length = max_completion_length
+        self.json_depth_limit = json_depth_limit
 
         self._verify_model()
         self.set_key(self.api_key)
@@ -94,7 +94,7 @@ def get_description(self) -> str:
 
         return self.supported_models()[self.model]
 
-    def model_output(self, response: Dict, max_completion_length = None) -> Dict:
+    def model_output(self, response: Dict, json_depth_limit = None) -> Dict:
 
         return [item["generated_text"] for item in response.json()]