codingutf8 2021 mesh tensorflow s t5 s and huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license flax mt5 model import jax numpy as jnp from utils import logging from t5 modelingflaxt5 import flaxt5encodermodel flaxt5forconditionalgeneration flaxt5model from configurationmt5 import mt5config logger logging getloggername configfordoc t5config copied from transformers models bart modelingflaxbart shifttokensright def shifttokensrightinputids jnp ndarray padtokenid int decoderstarttokenid int jnp ndarray shiftedinputids jnp zeroslikeinputids shiftedinputids shiftedinputids at 1 setinputids 1 shiftedinputids shiftedinputids at 0 setdecoderstarttokenid shiftedinputids jnp whereshiftedinputids 100 padtokenid shiftedinputids return shiftedinputids class flaxmt5modelflaxt5model r this class overrides flaxt5model please check the superclass for the appropriate documentation alongside usage examples examples python from transformers import flaxmt5model autotokenizer model flaxmt5model frompretrainedgooglemt5small tokenizer autotokenizer frompretrainedgooglemt5small article un offizier sagt dass weiter verhandelt werden muss in syrien summary weiter verhandlung in syrien inputs tokenizerarticle returntensorsnp decoderinputids tokenizertexttargetsummary returntensorsnp inputids outputs modelinputidsinputsinputids decoderinputidsdecoderinputids hiddenstates outputs lasthiddenstate modeltype mt5 configclass mt5config class flaxmt5encodermodelflaxt5encodermodel r this class overrides flaxt5encodermodel please check the superclass for the appropriate documentation alongside usage examples examples python from transformers import flaxt5encodermodel autotokenizer model flaxt5encodermodel frompretrainedgooglemt5small tokenizer autotokenizer frompretrainedgooglemt5small article un offizier sagt dass weiter verhandelt werden muss in syrien summary weiter verhandlung in syrien inputs tokenizerarticle returntensorsnp decoderinputids tokenizertexttargetsummary returntensorsnp inputids outputs modelinputidsinputsinputids hiddenstates outputs lasthiddenstate modeltype mt5 configclass mt5config class flaxmt5forconditionalgenerationflaxt5forconditionalgeneration r this class overrides flaxt5forconditionalgeneration please check the superclass for the appropriate documentation alongside usage examples examples python from transformers import flaxmt5forconditionalgeneration autotokenizer model flaxmt5forconditionalgeneration frompretrainedgooglemt5small tokenizer autotokenizer frompretrainedgooglemt5small article un offizier sagt dass weiter verhandelt werden muss in syrien summary weiter verhandlung in syrien inputs tokenizerarticle returntensorsnp decoderinputids tokenizertexttargetsummary returntensorsnp inputids outputs modelinputs decoderinputidsdecoderinputids logits outputs logits modeltype mt5 configclass mt5config coding utf 8 2021 mesh tensorflow s t5 s and huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license flax mt5 model copied from transformers models bart modeling_flax_bart shift_tokens_right shift input ids one token to the right this class overrides flaxt5model please check the superclass for the appropriate documentation alongside usage examples examples python from transformers import flaxmt5model autotokenizer model flaxmt5model from_pretrained google mt5 small tokenizer autotokenizer from_pretrained google mt5 small article un offizier sagt dass weiter verhandelt werden muss in syrien summary weiter verhandlung in syrien inputs tokenizer article return_tensors np decoder_input_ids tokenizer text_target summary return_tensors np input_ids outputs model input_ids inputs input_ids decoder_input_ids decoder_input_ids hidden_states outputs last_hidden_state this class overrides flaxt5encodermodel please check the superclass for the appropriate documentation alongside usage examples examples python from transformers import flaxt5encodermodel autotokenizer model flaxt5encodermodel from_pretrained google mt5 small tokenizer autotokenizer from_pretrained google mt5 small article un offizier sagt dass weiter verhandelt werden muss in syrien summary weiter verhandlung in syrien inputs tokenizer article return_tensors np decoder_input_ids tokenizer text_target summary return_tensors np input_ids outputs model input_ids inputs input_ids hidden_states outputs last_hidden_state this class overrides flaxt5forconditionalgeneration please check the superclass for the appropriate documentation alongside usage examples examples python from transformers import flaxmt5forconditionalgeneration autotokenizer model flaxmt5forconditionalgeneration from_pretrained google mt5 small tokenizer autotokenizer from_pretrained google mt5 small article un offizier sagt dass weiter verhandelt werden muss in syrien summary weiter verhandlung in syrien inputs tokenizer article return_tensors np decoder_input_ids tokenizer text_target summary return_tensors np input_ids outputs model inputs decoder_input_ids decoder_input_ids logits outputs logits
import jax.numpy as jnp

from ...utils import logging
from ..t5.modeling_flax_t5 import FlaxT5EncoderModel, FlaxT5ForConditionalGeneration, FlaxT5Model
from .configuration_mt5 import MT5Config


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "T5Config"



def shift_tokens_right(input_ids: jnp.ndarray, pad_token_id: int, decoder_start_token_id: int) -> jnp.ndarray:
    
    shifted_input_ids = jnp.zeros_like(input_ids)
    shifted_input_ids = shifted_input_ids.at[:, 1:].set(input_ids[:, :-1])
    shifted_input_ids = shifted_input_ids.at[:, 0].set(decoder_start_token_id)

    shifted_input_ids = jnp.where(shifted_input_ids == -100, pad_token_id, shifted_input_ids)
    return shifted_input_ids


class FlaxMT5Model(FlaxT5Model):
    r

    model_type = "mt5"
    config_class = MT5Config


class FlaxMT5EncoderModel(FlaxT5EncoderModel):
    r

    model_type = "mt5"
    config_class = MT5Config


class FlaxMT5ForConditionalGeneration(FlaxT5ForConditionalGeneration):
    r

    model_type = "mt5"
    config_class = MT5Config