WIP: add support for byt5 and mt5, closes #1655

Signed-off-by: sam-writer <sam.havens@writer.com>
NVIDIA · Dec 22, 2021 · d2fa90b · d2fa90b
1 parent 34a1675
commit d2fa90b
Show file tree

Hide file tree

Showing 2 changed files with 140 additions and 73 deletions.
diff --git a/demo/HuggingFace/T5/T5ModelConfig.py b/demo/HuggingFace/T5/T5ModelConfig.py
@@ -107,7 +107,7 @@ def __init__(self):
 
     def get_python_requirements(self):
         base_requirements = super().get_python_requirements()
-        base_requirements.append("transformers==4.6.1")
+        base_requirements.append("transformers>=4.8.0")
         return base_requirements
 
     def get_network_segments(self):
@@ -119,8 +119,8 @@ def get_network_segments(self):
         return T5ModelTRTConfig.NETWORK_SEGMENTS
 
     def get_metadata_string(self, metadata: NetworkMetadata) -> str:
-        # Remove redundant t5 name
-        metadata = metadata._replace(variant=metadata.variant.lstrip("t5-"))
+        # Remove redundant google/ if present
+        metadata = metadata._replace(variant=metadata.variant.lstrip("google/"))
         return super().get_metadata_string(metadata)
 
     @staticmethod