| 283 | ], |
| 284 | ) |
| 285 | def test_transformers_encoder(clear_transformers_cache, tmp_dir, device, model_name): |
| 286 | import torch |
| 287 | import transformers |
| 288 | |
| 289 | text = ["Hello world!", "Hello, my dog is cute"] |
| 290 | |
| 291 | model = transformers.AutoModel.from_pretrained(model_name) |
| 292 | tokenizer = transformers.AutoTokenizer.from_pretrained(model_name) |
| 293 | |
| 294 | inputs = tokenizer(text, return_tensors="pt", padding=True) |
| 295 | |
| 296 | inputs.to(device) |
| 297 | model.to(device) |
| 298 | |
| 299 | with torch.no_grad(): |
| 300 | outputs = model(**inputs) |
| 301 | |
| 302 | mask = inputs.attention_mask.unsqueeze(-1).cpu().numpy() |
| 303 | ref_last_hidden_state = outputs.last_hidden_state.cpu().numpy() |
| 304 | ref_pooler_output = ( |
| 305 | outputs.pooler_output.cpu().numpy() |
| 306 | if hasattr(outputs, "pooler_output") |
| 307 | else None |
| 308 | ) |
| 309 | |
| 310 | converter = ctranslate2.converters.TransformersConverter(model_name) |
| 311 | output_dir = str(tmp_dir.join("ctranslate2_model")) |
| 312 | output_dir = converter.convert(output_dir) |
| 313 | |
| 314 | encoder = ctranslate2.Encoder(output_dir, device=device) |
| 315 | |
| 316 | ids = [tokenizer(t).input_ids for t in text] |
| 317 | outputs = encoder.forward_batch(ids) |
| 318 | |
| 319 | last_hidden_state = _to_numpy(outputs.last_hidden_state, device) |
| 320 | assert last_hidden_state.shape == ref_last_hidden_state.shape |
| 321 | |
| 322 | last_hidden_state *= mask |
| 323 | ref_last_hidden_state *= mask |
| 324 | np.testing.assert_array_almost_equal( |
| 325 | last_hidden_state, ref_last_hidden_state, decimal=5 |
| 326 | ) |
| 327 | |
| 328 | if ref_pooler_output is not None: |
| 329 | pooler_output = _to_numpy(outputs.pooler_output, device) |
| 330 | assert pooler_output.shape == ref_pooler_output.shape |
| 331 | np.testing.assert_array_almost_equal( |
| 332 | pooler_output, ref_pooler_output, decimal=5 |
| 333 | ) |
| 334 | |
| 335 | |
| 336 | def _to_numpy(storage, device): |