Mistral nits (#57)

* fix(mistral): add missing sentencepiece dependency

It is required for Mistral models.

* feat(mistral): added bfloat16 dtyle by default

* chore(models): add warning about implicit conversion to bfloat16

optimum/tpu/modeling_mistral.py

@@ -1630,3 +1630,12 @@ def forward(
             hidden_states=outputs.hidden_states,
             attentions=outputs.attentions,
         )
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        # Unless specified otherwise, the model weights type will be bfloat16
+        if "torch_dtype" not in kwargs:
+            logger.warning_once("Defaulting to `torch_dtype=torch.bfloat16` for this model")
+        torch_dtype = kwargs.pop("torch_dtype", torch.bfloat16)
+        # forward to base implementation
+        return super().from_pretrained(pretrained_model_name_or_path, *model_args, torch_dtype=torch_dtype, **kwargs)

pyproject.toml

@@ -45,7 +45,8 @@ dependencies = [
     "transformers == 4.41.1",
     "torch >= 2.3.0, <= 2.4.0",
     "torch-xla[tpu] >= 2.3.0, <= 2.4.0",
-    "loguru == 0.6.0"
+    "loguru == 0.6.0",
+    "sentencepiece == 0.2.0",
 ]
 
 [tool.setuptools_scm]

text-generation-inference/server/pyproject.toml

@@ -17,6 +17,7 @@ dependencies = [
     'safetensors == 0.4.2',
     'transformers == 4.41.1',
     'loguru == 0.6.0',
+    "sentencepiece == 0.2.0",
 ]
 
 [tool.setuptools]