Merge pull request #13 from Genaios/fix/provider-bugs

Add mistral to bedrock and upgrade anthropic

text_machina/src/constrainers/length.py

@@ -45,20 +45,19 @@ def get_constraints(self) -> Dict[str, int]:
         """
         min_new_tokens, max_new_tokens = self.estimate()
 
-        if self.provider in ["openai", "azure_openai"]:
-            return {"max_tokens": max_new_tokens}
-        elif self.provider == "anthropic":
-            return {"max_tokens_to_sample": max_new_tokens}
-        elif self.provider == "cohere":
+        if self.provider in [
+            "openai",
+            "azure_openai",
+            "anthropic",
+            "cohere",
+            "ai21",
+            "inference_server",
+        ]:
             return {"max_tokens": max_new_tokens}
         elif self.provider == "vertex":
             return {"max_output_tokens": max_new_tokens}
         elif self.provider == "bedrock":
             return {"maxTokenCount": max_new_tokens}
-        elif self.provider == "ai21":
-            return {"max_tokens": max_new_tokens}
-        elif self.provider == "inference_server":
-            return {"max_tokens": max_new_tokens}
         return {
             "max_new_tokens": max_new_tokens,
             "min_new_tokens": min_new_tokens,

text_machina/src/models/anthropic.py

@@ -34,11 +34,15 @@ def generate_completion(
         generation_config: Dict,
     ) -> str:
         try:
-            completion = self.client.completions.create(
-                model=self.model_config.model_name,
-                prompt=prompt,
-                **generation_config,
-            ).completion
+            completion = (
+                self.client.messages.create(
+                    model=self.model_config.model_name,
+                    messages=[{"role": "user", "content": prompt}],
+                    **generation_config,
+                )
+                .content[0]
+                .text
+            )
         except Exception as e:
             _logger.info(f"Unrecoverable exception during the request: {e}")
             return GENERATION_ERROR

text_machina/src/models/bedrock.py

@@ -85,6 +85,7 @@ def get_request_body(
             "anthropic",
             "cohere",
             "meta",
+            "mistral",
         }
 
         if bedrock_provider == "amazon":
@@ -118,6 +119,13 @@ def get_request_body(
                     "maxTokenCount"
                 )
             request_body = {"prompt": prompt, **generation_config}
+        elif bedrock_provider == "mistral":
+            if "maxTokenCount" in generation_config:
+                generation_config["max_tokens"] = generation_config.pop(
+                    "maxTokenCount"
+                )
+            request_body = {"prompt": prompt, **generation_config}
+
         return json.dumps(request_body)
 
     def get_completion_from_response_body(self, response_body: Dict) -> str:
@@ -138,6 +146,7 @@ def get_completion_from_response_body(self, response_body: Dict) -> str:
             "anthropic",
             "cohere",
             "meta",
+            "mistral",
         }
 
         if bedrock_provider == "amazon":
@@ -150,5 +159,7 @@ def get_completion_from_response_body(self, response_body: Dict) -> str:
             completion = response_body["generations"][0]["text"]
         elif bedrock_provider == "meta":
             completion = response_body["generation"]
+        elif bedrock_provider == "mistral":
+            completion = response_body["outputs"][0]["text"]
 
         return completion

text_machina/version.py

@@ -1,10 +1,12 @@
 _MAJOR = "0"
 _MINOR = "2"
-_REVISION = "7"
+_REVISION = "8"
 
 VERSION_SHORT = "{0}.{1}".format(_MAJOR, _MINOR)
 VERSION = "{0}.{1}.{2}".format(_MAJOR, _MINOR, _REVISION)
 
+__version__ = VERSION
+
 
 def _is_newer_than(version: str) -> bool:
     """True if current version is newer than 'version'."""