fix(server): Fixing T5 in case the names are mixed up. (#475)

2023-06-20 18:03:36 +02:00 · 2023-06-20 18:03:36 +02:00 · c9c65ab323
parent 53aa9194c8
commit c9c65ab323
1 changed files with 4 additions and 1 deletions
--- a/server/text_generation_server/models/custom_modeling/t5_modeling.py
+++ b/server/text_generation_server/models/custom_modeling/t5_modeling.py
@ -1001,7 +1001,10 @@ class T5ForConditionalGeneration(T5PreTrainedModel):
        super().__init__(config)
        self.model_dim = config.d_model

+        try:
            self.shared = TensorParallelEmbedding(prefix="shared", weights=weights)
+        except RuntimeError:
+            self.shared = TensorParallelEmbedding(prefix="encoder.embed_tokens", weights=weights)

        encoder_config = copy.deepcopy(config)
        encoder_config.is_decoder = False