(
self,
model_dir="google/t5-v1_1-xxl",
device="cuda",
max_length=77,
freeze=True,
cache_dir=None,
)
| 187 | """Uses the T5 transformer encoder for text""" |
| 188 | |
| 189 | def __init__( |
| 190 | self, |
| 191 | model_dir="google/t5-v1_1-xxl", |
| 192 | device="cuda", |
| 193 | max_length=77, |
| 194 | freeze=True, |
| 195 | cache_dir=None, |
| 196 | ): |
| 197 | super().__init__() |
| 198 | if model_dir is not "google/t5-v1_1-xxl": |
| 199 | self.tokenizer = T5Tokenizer.from_pretrained(model_dir) |
| 200 | self.transformer = T5EncoderModel.from_pretrained(model_dir) |
| 201 | else: |
| 202 | self.tokenizer = T5Tokenizer.from_pretrained(model_dir, cache_dir=cache_dir) |
| 203 | self.transformer = T5EncoderModel.from_pretrained(model_dir, cache_dir=cache_dir) |
| 204 | self.device = device |
| 205 | self.max_length = max_length |
| 206 | if freeze: |
| 207 | self.freeze() |
| 208 | |
| 209 | def freeze(self): |
| 210 | self.transformer = self.transformer.eval() |
nothing calls this directly
no test coverage detected