* @param {PretrainedConfig} config * @param {{ prefix?: string }} [options] * @returns {Set }
(config, { prefix = 'past_key_values' } = {})
| 417 | * @returns {Set<string>} |
| 418 | */ |
| 419 | function getKeyValueNames(config, { prefix = 'past_key_values' } = {}) { |
| 420 | /** @type {Set<string>} */ |
| 421 | const names = new Set(); |
| 422 | const normalized_config = config.normalized_config; |
| 423 | |
| 424 | if ( |
| 425 | normalized_config.is_encoder_decoder && |
| 426 | 'num_encoder_heads' in normalized_config && |
| 427 | 'num_decoder_heads' in normalized_config |
| 428 | ) { |
| 429 | for (let i = 0; i < normalized_config.num_decoder_layers; ++i) { |
| 430 | names.add(`${prefix}.${i}.encoder.key`); |
| 431 | names.add(`${prefix}.${i}.encoder.value`); |
| 432 | names.add(`${prefix}.${i}.decoder.key`); |
| 433 | names.add(`${prefix}.${i}.decoder.value`); |
| 434 | } |
| 435 | } else if (normalized_config.multi_query) { |
| 436 | // e.g., for `gpt_bigcode` |
| 437 | for (let i = 0; i < normalized_config.num_layers; ++i) { |
| 438 | names.add(`${prefix}.${i}.key_value`); |
| 439 | } |
| 440 | } else { |
| 441 | for (let i = 0; i < normalized_config.num_layers; ++i) { |
| 442 | names.add(`${prefix}.${i}.key`); |
| 443 | names.add(`${prefix}.${i}.value`); |
| 444 | } |
| 445 | } |
| 446 | |
| 447 | return names; |
| 448 | } |
| 449 | /** |
| 450 | * Base class for all configuration classes. For more information, see the corresponding |
| 451 | * [Python documentation](https://huggingface.co/docs/transformers/main/en/main_classes/configuration#transformers.PretrainedConfig). |