From 3e8d227247abb32bbf5cd76308cb514c66b7906c Mon Sep 17 00:00:00 2001 From: Joshua Lochner Date: Thu, 9 Nov 2023 16:43:51 +0200 Subject: [PATCH] Add support for Falcon and Mistral models (#379) * By default, do not add special tokens in text-generation See https://github.com/huggingface/transformers/blob/147e8ce4ae23f6f360ac0bc6e162cdf9995516b7/src/transformers/pipelines/text_generation.py#L106 * Add support for mistral models * Add support for Falcon models * Replace `batch_size` with variable * Add Falcon to list of supported models * Fix typing issue with bigint literals --- README.md | 2 + docs/snippets/6_supported-models.snippet | 2 + scripts/convert.py | 8 ++ scripts/supported_models.py | 9 ++ src/models.js | 139 +++++++++++++++++++++-- src/pipelines.js | 4 + 6 files changed, 156 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 2933a31e8..7ee4e5a09 100644 --- a/README.md +++ b/README.md @@ -273,6 +273,7 @@ You can refine your search by selecting the task you're interested in (e.g., [te 1. **[DETR](https://huggingface.co/docs/transformers/model_doc/detr)** (from Facebook) released with the paper [End-to-End Object Detection with Transformers](https://arxiv.org/abs/2005.12872) by Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, Sergey Zagoruyko. 1. **[DistilBERT](https://huggingface.co/docs/transformers/model_doc/distilbert)** (from HuggingFace), released together with the paper [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://arxiv.org/abs/1910.01108) by Victor Sanh, Lysandre Debut and Thomas Wolf. The same method has been applied to compress GPT2 into [DistilGPT2](https://github.com/huggingface/transformers/tree/main/examples/research_projects/distillation), RoBERTa into [DistilRoBERTa](https://github.com/huggingface/transformers/tree/main/examples/research_projects/distillation), Multilingual BERT into [DistilmBERT](https://github.com/huggingface/transformers/tree/main/examples/research_projects/distillation) and a German version of DistilBERT. 1. **[Donut](https://huggingface.co/docs/transformers/model_doc/donut)** (from NAVER), released together with the paper [OCR-free Document Understanding Transformer](https://arxiv.org/abs/2111.15664) by Geewook Kim, Teakgyu Hong, Moonbin Yim, Jeongyeon Nam, Jinyoung Park, Jinyeong Yim, Wonseok Hwang, Sangdoo Yun, Dongyoon Han, Seunghyun Park. +1. **[Falcon](https://huggingface.co/docs/transformers/model_doc/falcon)** (from Technology Innovation Institute) by Almazrouei, Ebtesam and Alobeidli, Hamza and Alshamsi, Abdulaziz and Cappelli, Alessandro and Cojocaru, Ruxandra and Debbah, Merouane and Goffinet, Etienne and Heslow, Daniel and Launay, Julien and Malartic, Quentin and Noune, Badreddine and Pannier, Baptiste and Penedo, Guilherme. 1. **[FLAN-T5](https://huggingface.co/docs/transformers/model_doc/flan-t5)** (from Google AI) released in the repository [google-research/t5x](https://github.com/google-research/t5x/blob/main/docs/models.md#flan-t5-checkpoints) by Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Eric Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, Albert Webson, Shixiang Shane Gu, Zhuyun Dai, Mirac Suzgun, Xinyun Chen, Aakanksha Chowdhery, Sharan Narang, Gaurav Mishra, Adams Yu, Vincent Zhao, Yanping Huang, Andrew Dai, Hongkun Yu, Slav Petrov, Ed H. Chi, Jeff Dean, Jacob Devlin, Adam Roberts, Denny Zhou, Quoc V. Le, and Jason Wei 1. **[GPT Neo](https://huggingface.co/docs/transformers/model_doc/gpt_neo)** (from EleutherAI) released in the repository [EleutherAI/gpt-neo](https://github.com/EleutherAI/gpt-neo) by Sid Black, Stella Biderman, Leo Gao, Phil Wang and Connor Leahy. 1. **[GPT NeoX](https://huggingface.co/docs/transformers/model_doc/gpt_neox)** (from EleutherAI) released with the paper [GPT-NeoX-20B: An Open-Source Autoregressive Language Model](https://arxiv.org/abs/2204.06745) by Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, Michael Pieler, USVSN Sai Prashanth, Shivanshu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, Samuel Weinbach @@ -287,6 +288,7 @@ You can refine your search by selecting the task you're interested in (e.g., [te 1. **[MarianMT](https://huggingface.co/docs/transformers/model_doc/marian)** Machine translation models trained using [OPUS](http://opus.nlpl.eu/) data by Jörg Tiedemann. The [Marian Framework](https://marian-nmt.github.io/) is being developed by the Microsoft Translator Team. 1. **[mBART](https://huggingface.co/docs/transformers/model_doc/mbart)** (from Facebook) released with the paper [Multilingual Denoising Pre-training for Neural Machine Translation](https://arxiv.org/abs/2001.08210) by Yinhan Liu, Jiatao Gu, Naman Goyal, Xian Li, Sergey Edunov, Marjan Ghazvininejad, Mike Lewis, Luke Zettlemoyer. 1. **[mBART-50](https://huggingface.co/docs/transformers/model_doc/mbart)** (from Facebook) released with the paper [Multilingual Translation with Extensible Multilingual Pretraining and Finetuning](https://arxiv.org/abs/2008.00401) by Yuqing Tang, Chau Tran, Xian Li, Peng-Jen Chen, Naman Goyal, Vishrav Chaudhary, Jiatao Gu, Angela Fan. +1. **[Mistral](https://huggingface.co/docs/transformers/model_doc/mistral)** (from Mistral AI) by The [Mistral AI](https://mistral.ai) team: Albert Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, William El Sayed. 1. **[MMS](https://huggingface.co/docs/transformers/model_doc/mms)** (from Facebook) released with the paper [Scaling Speech Technology to 1,000+ Languages](https://arxiv.org/abs/2305.13516) by Vineel Pratap, Andros Tjandra, Bowen Shi, Paden Tomasello, Arun Babu, Sayani Kundu, Ali Elkahky, Zhaoheng Ni, Apoorv Vyas, Maryam Fazel-Zarandi, Alexei Baevski, Yossi Adi, Xiaohui Zhang, Wei-Ning Hsu, Alexis Conneau, Michael Auli. 1. **[MobileBERT](https://huggingface.co/docs/transformers/model_doc/mobilebert)** (from CMU/Google Brain) released with the paper [MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices](https://arxiv.org/abs/2004.02984) by Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. 1. **[MobileViT](https://huggingface.co/docs/transformers/model_doc/mobilevit)** (from Apple) released with the paper [MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer](https://arxiv.org/abs/2110.02178) by Sachin Mehta and Mohammad Rastegari. diff --git a/docs/snippets/6_supported-models.snippet b/docs/snippets/6_supported-models.snippet index b35b03496..3b7315410 100644 --- a/docs/snippets/6_supported-models.snippet +++ b/docs/snippets/6_supported-models.snippet @@ -18,6 +18,7 @@ 1. **[DETR](https://huggingface.co/docs/transformers/model_doc/detr)** (from Facebook) released with the paper [End-to-End Object Detection with Transformers](https://arxiv.org/abs/2005.12872) by Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, Sergey Zagoruyko. 1. **[DistilBERT](https://huggingface.co/docs/transformers/model_doc/distilbert)** (from HuggingFace), released together with the paper [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://arxiv.org/abs/1910.01108) by Victor Sanh, Lysandre Debut and Thomas Wolf. The same method has been applied to compress GPT2 into [DistilGPT2](https://github.com/huggingface/transformers/tree/main/examples/research_projects/distillation), RoBERTa into [DistilRoBERTa](https://github.com/huggingface/transformers/tree/main/examples/research_projects/distillation), Multilingual BERT into [DistilmBERT](https://github.com/huggingface/transformers/tree/main/examples/research_projects/distillation) and a German version of DistilBERT. 1. **[Donut](https://huggingface.co/docs/transformers/model_doc/donut)** (from NAVER), released together with the paper [OCR-free Document Understanding Transformer](https://arxiv.org/abs/2111.15664) by Geewook Kim, Teakgyu Hong, Moonbin Yim, Jeongyeon Nam, Jinyoung Park, Jinyeong Yim, Wonseok Hwang, Sangdoo Yun, Dongyoon Han, Seunghyun Park. +1. **[Falcon](https://huggingface.co/docs/transformers/model_doc/falcon)** (from Technology Innovation Institute) by Almazrouei, Ebtesam and Alobeidli, Hamza and Alshamsi, Abdulaziz and Cappelli, Alessandro and Cojocaru, Ruxandra and Debbah, Merouane and Goffinet, Etienne and Heslow, Daniel and Launay, Julien and Malartic, Quentin and Noune, Badreddine and Pannier, Baptiste and Penedo, Guilherme. 1. **[FLAN-T5](https://huggingface.co/docs/transformers/model_doc/flan-t5)** (from Google AI) released in the repository [google-research/t5x](https://github.com/google-research/t5x/blob/main/docs/models.md#flan-t5-checkpoints) by Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Eric Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, Albert Webson, Shixiang Shane Gu, Zhuyun Dai, Mirac Suzgun, Xinyun Chen, Aakanksha Chowdhery, Sharan Narang, Gaurav Mishra, Adams Yu, Vincent Zhao, Yanping Huang, Andrew Dai, Hongkun Yu, Slav Petrov, Ed H. Chi, Jeff Dean, Jacob Devlin, Adam Roberts, Denny Zhou, Quoc V. Le, and Jason Wei 1. **[GPT Neo](https://huggingface.co/docs/transformers/model_doc/gpt_neo)** (from EleutherAI) released in the repository [EleutherAI/gpt-neo](https://github.com/EleutherAI/gpt-neo) by Sid Black, Stella Biderman, Leo Gao, Phil Wang and Connor Leahy. 1. **[GPT NeoX](https://huggingface.co/docs/transformers/model_doc/gpt_neox)** (from EleutherAI) released with the paper [GPT-NeoX-20B: An Open-Source Autoregressive Language Model](https://arxiv.org/abs/2204.06745) by Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Horace He, Connor Leahy, Kyle McDonell, Jason Phang, Michael Pieler, USVSN Sai Prashanth, Shivanshu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, Samuel Weinbach @@ -32,6 +33,7 @@ 1. **[MarianMT](https://huggingface.co/docs/transformers/model_doc/marian)** Machine translation models trained using [OPUS](http://opus.nlpl.eu/) data by Jörg Tiedemann. The [Marian Framework](https://marian-nmt.github.io/) is being developed by the Microsoft Translator Team. 1. **[mBART](https://huggingface.co/docs/transformers/model_doc/mbart)** (from Facebook) released with the paper [Multilingual Denoising Pre-training for Neural Machine Translation](https://arxiv.org/abs/2001.08210) by Yinhan Liu, Jiatao Gu, Naman Goyal, Xian Li, Sergey Edunov, Marjan Ghazvininejad, Mike Lewis, Luke Zettlemoyer. 1. **[mBART-50](https://huggingface.co/docs/transformers/model_doc/mbart)** (from Facebook) released with the paper [Multilingual Translation with Extensible Multilingual Pretraining and Finetuning](https://arxiv.org/abs/2008.00401) by Yuqing Tang, Chau Tran, Xian Li, Peng-Jen Chen, Naman Goyal, Vishrav Chaudhary, Jiatao Gu, Angela Fan. +1. **[Mistral](https://huggingface.co/docs/transformers/model_doc/mistral)** (from Mistral AI) by The [Mistral AI](https://mistral.ai) team: Albert Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, William El Sayed. 1. **[MMS](https://huggingface.co/docs/transformers/model_doc/mms)** (from Facebook) released with the paper [Scaling Speech Technology to 1,000+ Languages](https://arxiv.org/abs/2305.13516) by Vineel Pratap, Andros Tjandra, Bowen Shi, Paden Tomasello, Arun Babu, Sayani Kundu, Ali Elkahky, Zhaoheng Ni, Apoorv Vyas, Maryam Fazel-Zarandi, Alexei Baevski, Yossi Adi, Xiaohui Zhang, Wei-Ning Hsu, Alexis Conneau, Michael Auli. 1. **[MobileBERT](https://huggingface.co/docs/transformers/model_doc/mobilebert)** (from CMU/Google Brain) released with the paper [MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices](https://arxiv.org/abs/2004.02984) by Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. 1. **[MobileViT](https://huggingface.co/docs/transformers/model_doc/mobilevit)** (from Apple) released with the paper [MobileViT: Light-weight, General-purpose, and Mobile-friendly Vision Transformer](https://arxiv.org/abs/2110.02178) by Sachin Mehta and Mohammad Rastegari. diff --git a/scripts/convert.py b/scripts/convert.py index 0a68cd5ca..ffc999bfa 100644 --- a/scripts/convert.py +++ b/scripts/convert.py @@ -67,6 +67,14 @@ 'per_channel': False, 'reduce_range': False, }, + 'mistral': { + 'per_channel': False, + 'reduce_range': False, + }, + 'falcon': { + 'per_channel': False, + 'reduce_range': False, + }, # Encoder-decoder models 'whisper': { diff --git a/scripts/supported_models.py b/scripts/supported_models.py index 5d74b1bce..f59122ae9 100644 --- a/scripts/supported_models.py +++ b/scripts/supported_models.py @@ -208,6 +208,11 @@ # Document Question Answering 'naver-clova-ix/donut-base-finetuned-docvqa', ], + 'falcon': [ + # Text generation + 'Rocketknight1/tiny-random-falcon-7b', + 'fxmarty/really-tiny-falcon-testing', + ], 'gpt_neo': [ # Text generation 'EleutherAI/gpt-neo-125M', @@ -303,6 +308,10 @@ 'facebook/mbart-large-50-many-to-one-mmt', 'facebook/mbart-large-50', ], + 'mistral': [ + # Text generation + 'echarlaix/tiny-random-mistral', + ], 'mobilebert': [ # Zero-shot classification 'typeform/mobilebert-uncased-mnli', diff --git a/src/models.js b/src/models.js index 57c8a8cf8..b2e918103 100644 --- a/src/models.js +++ b/src/models.js @@ -283,6 +283,41 @@ function prepareAttentionMask(self, tokens) { } } +/** + * Add position IDs to the feeds object. + * @param {Object} session The inference session. + * @param {Object} feeds The input to the model. + * @param {boolean} use_cache_branch Whether to use the cache branch of the model. + * @returns {void} + * @private + */ +function preparePositionIds(session, feeds, use_cache_branch) { + if (!session.inputNames.includes('position_ids')) return; + + const data = new BigInt64Array(feeds.attention_mask.data.length); + + // Compute cumulative sum of the attention mask along the sequence length dimension + for (let i = 0; i < feeds.attention_mask.dims[0]; ++i) { + let start = i * feeds.attention_mask.dims[1]; + let sum = BigInt(0); + for (let j = 0; j < feeds.attention_mask.dims[1]; ++j) { + const index = start + j; + if (feeds.attention_mask.data[index] === 0n) { + data[index] = BigInt(1); + } else { // === 1n + data[index] = sum; + sum += feeds.attention_mask.data[index]; + } + } + } + + feeds.position_ids = new Tensor('int64', data, feeds.attention_mask.dims); + + if (use_cache_branch) { + feeds.position_ids = feeds.position_ids.slice(null, -1).unsqueeze_(-1); + } +} + /** * Creates a boolean tensor with a single value. * @param {boolean} value The value of the tensor. @@ -312,12 +347,18 @@ async function seq2seqForward(self, model_inputs) { let decoderFeeds = { input_ids: model_inputs.decoder_input_ids, encoder_hidden_states: encoder_outputs, - use_cache_branch: boolTensor(!!past_key_values) }; + const use_cache_branch = !!past_key_values; + + if (self.decoder_merged_session.inputNames.includes('use_cache_branch')) { + decoderFeeds.use_cache_branch = boolTensor(use_cache_branch); + } if (self.decoder_merged_session.inputNames.includes('encoder_attention_mask')) { decoderFeeds.encoder_attention_mask = model_inputs.attention_mask } + + preparePositionIds(self.decoder_merged_session, decoderFeeds, use_cache_branch); self.addPastKeyValues(decoderFeeds, past_key_values); const decoderResults = await sessionRun(self.decoder_merged_session, decoderFeeds); @@ -467,8 +508,14 @@ async function decoderForward(self, model_inputs) { let decoderFeeds = { input_ids: input_ids, attention_mask: attention_mask ?? prepareAttentionMask(self, input_ids), - use_cache_branch: boolTensor(!!past_key_values) } + const use_cache_branch = !!past_key_values; + + if (self.session.inputNames.includes('use_cache_branch')) { + decoderFeeds.use_cache_branch = boolTensor(use_cache_branch); + } + + preparePositionIds(self.session, decoderFeeds, use_cache_branch); self.addPastKeyValues(decoderFeeds, past_key_values); @@ -1216,12 +1263,14 @@ export class PreTrainedModel extends Callable { Object.assign(decoderFeeds, pastKeyValues) } else { // TODO support batches (i.e., batch_size > 1) + const batch_size = 1; + // @ts-ignore if (this.config.is_encoder_decoder && (this.add_encoder_pkv ?? true)) { // @ts-ignore - let encoder_dims = [1, this.num_encoder_heads, 0, this.encoder_dim_kv]; + let encoder_dims = [batch_size, this.num_encoder_heads, 0, this.encoder_dim_kv]; // @ts-ignore - let decoder_dims = [1, this.num_decoder_heads, 0, this.decoder_dim_kv]; + let decoder_dims = [batch_size, this.num_decoder_heads, 0, this.decoder_dim_kv]; // @ts-ignore for (let i = 0; i < this.num_decoder_layers; ++i) { decoderFeeds[`past_key_values.${i}.encoder.key`] = new Tensor('float32', [], encoder_dims) @@ -1229,9 +1278,18 @@ export class PreTrainedModel extends Callable { decoderFeeds[`past_key_values.${i}.decoder.key`] = new Tensor('float32', [], decoder_dims) decoderFeeds[`past_key_values.${i}.decoder.value`] = new Tensor('float32', [], decoder_dims) } + } else if (this.config.model_type === 'falcon') { + // NOTE: Custom implementation for Falcon + // @ts-ignore + let dims = [batch_size * this.num_heads, 0, this.dim_kv] + // @ts-ignore + for (let i = 0; i < this.num_layers; ++i) { + decoderFeeds[`past_key_values.${i}.key`] = new Tensor('float32', [], dims) + decoderFeeds[`past_key_values.${i}.value`] = new Tensor('float32', [], dims) + } } else if (this.config.multi_query) { // e.g., for `gpt_bigcode` // @ts-ignore - let dims = [1, 0, 2 * this.dim_kv] + let dims = [batch_size * this.num_heads, 0, 2 * this.dim_kv] // @ts-ignore for (let i = 0; i < this.num_layers; ++i) { decoderFeeds[`past_key_values.${i}.key_value`] = new Tensor('float32', [], dims) @@ -1240,9 +1298,9 @@ export class PreTrainedModel extends Callable { // NOTE: Custom implementation for Bloom // @ts-ignore - let keyDims = [1 * this.num_heads, this.dim_kv, 0] // [batch_size x num_heads,64,past_sequence_length] + let keyDims = [batch_size * this.num_heads, this.dim_kv, 0] // [batch_size x num_heads,64,past_sequence_length] // @ts-ignore - let valueDims = [1 * this.num_heads, 0, this.dim_kv] // [batch_size x num_heads,past_sequence_length,64] + let valueDims = [batch_size * this.num_heads, 0, this.dim_kv] // [batch_size x num_heads,past_sequence_length,64] // @ts-ignore for (let i = 0; i < this.num_layers; ++i) { decoderFeeds[`past_key_values.${i}.key`] = new Tensor('float32', [], keyDims) @@ -1250,7 +1308,7 @@ export class PreTrainedModel extends Callable { } } else { // Decoder-only // @ts-ignore - let dims = [1, this.num_heads, 0, this.dim_kv] + let dims = [batch_size, this.num_heads, 0, this.dim_kv] // @ts-ignore for (let i = 0; i < this.num_layers; ++i) { decoderFeeds[`past_key_values.${i}.key`] = new Tensor('float32', [], dims) @@ -3738,6 +3796,67 @@ export class TrOCRForCausalLM extends TrOCRPreTrainedModel { } ////////////////////////////////////////////////// +////////////////////////////////////////////////// +// Mistral models +/** + * The bare Mistral Model outputting raw hidden-states without any specific head on top. + */ +export class MistralPreTrainedModel extends PreTrainedModel { + /** + * Creates a new instance of the `MistralPreTrainedModel` class. + * @param {Object} config The configuration of the model. + * @param {any} session The ONNX session containing the model weights. + * @param {GenerationConfig} generation_config The generation configuration. + */ + constructor(config, session, generation_config) { + super(config, session); + this.generation_config = generation_config; + + // config doesn't contain pad_token_id, so we assume it is the eos_token_id + this.config.pad_token_id = this.config.eos_token_id + + this.num_heads = this.config.num_key_value_heads; + this.num_layers = this.config.num_hidden_layers; + this.dim_kv = this.config.hidden_size / this.config.num_attention_heads; + } +} + +export class MistralModel extends MistralPreTrainedModel { } + +export class MistralForCausalLM extends MistralPreTrainedModel { } +////////////////////////////////////////////////// + +////////////////////////////////////////////////// +// Falcon models +/** + * The bare Falcon Model outputting raw hidden-states without any specific head on top. + */ +export class FalconPreTrainedModel extends PreTrainedModel { + /** + * Creates a new instance of the `FalconPreTrainedModel` class. + * @param {Object} config The configuration of the model. + * @param {any} session The ONNX session containing the model weights. + * @param {GenerationConfig} generation_config The generation configuration. + */ + constructor(config, session, generation_config) { + super(config, session); + this.generation_config = generation_config; + + // config doesn't contain pad_token_id, so we assume it is the eos_token_id + this.config.pad_token_id = this.config.eos_token_id + + this.num_heads = this.config.num_attention_heads; + this.num_layers = this.config.num_hidden_layers; + this.dim_kv = this.config.hidden_size / this.config.num_attention_heads; + } +} + +export class FalconModel extends FalconPreTrainedModel { } + +export class FalconForCausalLM extends FalconPreTrainedModel { } +////////////////////////////////////////////////// + + ////////////////////////////////////////////////// // AutoModels, used to simplify construction of PreTrainedModels // (uses config to instantiate correct class) @@ -3864,6 +3983,8 @@ const MODEL_MAPPING_NAMES_DECODER_ONLY = new Map([ ['llama', ['LlamaModel', LlamaModel]], ['mpt', ['MptModel', MptModel]], ['opt', ['OPTModel', OPTModel]], + ['mistral', ['MistralModel', MistralModel]], + ['falcon', ['FalconModel', FalconModel]], ]); const MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES = new Map([ @@ -3928,6 +4049,8 @@ const MODEL_WITH_LM_HEAD_MAPPING_NAMES = new Map([ ['mpt', ['MptForCausalLM', MptForCausalLM]], ['opt', ['OPTForCausalLM', OPTForCausalLM]], ['mbart', ['MBartForCausalLM', MBartForCausalLM]], + ['mistral', ['MistralForCausalLM', MistralForCausalLM]], + ['falcon', ['FalconForCausalLM', FalconForCausalLM]], ['trocr', ['TrOCRForCausalLM', TrOCRForCausalLM]], ]); diff --git a/src/pipelines.js b/src/pipelines.js index d40b8929c..8842b6337 100644 --- a/src/pipelines.js +++ b/src/pipelines.js @@ -686,8 +686,12 @@ export class TextGenerationPipeline extends Pipeline { texts = [texts]; } + // By default, do not add special tokens + const add_special_tokens = generate_kwargs.add_special_tokens ?? false; + this.tokenizer.padding_side = 'left'; let inputs = this.tokenizer(texts, { + add_special_tokens, padding: true, truncation: true, });