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

c9c65ab3 · Nicolas Patry · GitHub · 53aa9194 · c9c65ab3
Unverified Commit c9c65ab3 authored Jun 20, 2023 by Nicolas Patry Committed by GitHub Jun 20, 2023
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

server/text_generation_server/models/custom_modeling/t5_modeling.py ...t_generation_server/models/custom_modeling/t5_modeling.py +4 -1

No files found.
--- 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