v0.8.2

Author: OlivierDehaene

Cargo.lock

@@ -8,7 +8,7 @@ members = [
 ]
 
 [workspace.package]
-version = "0.8.1"
+version = "0.8.2"
 edition = "2021"
 authors = ["Olivier Dehaene"]
 homepage = "https://github.com/huggingface/text-generation-inference"

Cargo.toml

@@ -10,7 +10,7 @@
       "name": "Apache 2.0",
       "url": "https://www.apache.org/licenses/LICENSE-2.0"
     },
-    "version": "0.8.1"
+    "version": "0.8.2"
   },
   "paths": {
     "/": {

docs/openapi.json

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "text-generation-server"
-version = "0.8.1"
+version = "0.8.2"
 description = "Text Generation Inference Python gRPC Server"
 authors = ["Olivier Dehaene <[email protected]>"]
 

server/pyproject.toml

@@ -138,7 +138,9 @@ def get_model(
                 trust_remote_code=trust_remote_code,
             )
 
-    config_dict, _ = PretrainedConfig.get_config_dict(model_id, revision=revision, trust_remote_code=trust_remote_code)
+    config_dict, _ = PretrainedConfig.get_config_dict(
+        model_id, revision=revision, trust_remote_code=trust_remote_code
+    )
     model_type = config_dict["model_type"]
 
     if model_type == "gpt_bigcode":

server/text_generation_server/models/__init__.py

@@ -87,7 +87,9 @@ def load_weights(
     ):
         for filename in filenames:
             with safe_open(
-                    filename, framework="pt", device=str(device) if quantize is None else "cpu"
+                filename,
+                framework="pt",
+                device=str(device) if quantize is None else "cpu",
             ) as f:
                 for key in f.keys():
                     value = f.get_tensor(key)
@@ -148,11 +150,13 @@ def load_weights(
                             module._parameters[param_name][: value.shape[0]] = value
                         elif "kv_attn.weight" in key:
                             module._parameters[param_name][
-                                model.transformer.head_size * model.transformer.num_heads :
+                                model.transformer.head_size
+                                * model.transformer.num_heads :
                             ] = value
                         elif "kv_attn.bias" in key:
                             module._parameters[param_name][
-                                model.transformer.head_size * model.transformer.num_heads :
+                                model.transformer.head_size
+                                * model.transformer.num_heads :
                             ] = value
                         else:
                             if current_parameter_tensor.shape != value.shape: