Package: keras-transformer
Keras Transformer[中文|English]Implementation oftransformerfor seq2seq tasks.Installpipinstallkeras-transformerUsageTrainimportnumpyasnpfromkeras_transformerimportget_model# Build a small toy token dictionarytokens='all work and no play makes jack a dull boy'.split(' ')token_dict={'<PAD>':0,'<START>':1,'<END>':2,}fortokenintokens:iftokennotintoken_dict:token_dict[token]=len(token_dict)# Generate toy dataencoder_inputs_no_padding=[]encoder_inputs,decoder_inputs,decoder_outputs=[],[],[]foriinrange(1,len(tokens)-1):encode_tokens,decode_tokens=tokens[:i],tokens[i:]encode_tokens=['<START>']+encode_tokens+['<END>']+['<PAD>']*(len(tokens)-len(encode_tokens))output_tokens=decode_tokens+['<END>','<PAD>']+['<PAD>']*(len(tokens)-len(decode_tokens))decode_tokens=['<START>']+decode_tokens+['<END>']+['<PAD>']*(len(tokens)-len(decode_tokens))encode_tokens=list(map(lambdax:token_dict[x],encode_tokens))decode_tokens=list(map(lambdax:token_dict[x],decode_tokens))output_tokens=list(map(lambdax:[token_dict[x]],output_tokens))encoder_inputs_no_padding.append(encode_tokens[:i+2])encoder_inputs.append(encode_tokens)decoder_inputs.append(decode_tokens)decoder_outputs.append(output_tokens)# Build the modelmodel=get_model(token_num=len(token_dict),embed_dim=30,encoder_num=3,decoder_num=2,head_num=3,hidden_dim=120,attention_activation='relu',feed_forward_activation='relu',dropout_rate=0.05,embed_weights=np.random.random((13,30)),)model.compile(optimizer='adam',loss='sparse_categorical_crossentropy',)model.summary()# Train the modelmodel.fit(x=[np.asarray(encoder_inputs*1000),np.asarray(decoder_inputs*1000)],y=np.asarray(decoder_outputs*1000),epochs=5,)Predictfromkeras_transformerimportdecodedecoded=decode(model,encoder_inputs_no_padding,start_token=token_dict['<START>'],end_token=token_dict['<END>'],pad_token=token_dict['<PAD>'],max_len=100,)token_dict_rev={v:kfork,vintoken_dict.items()}foriinrange(len(decoded)):print(' '.join(map(lambdax:token_dict_rev[x],decoded[i][1:-1])))Translationimportnumpyasnpfromkeras_transformerimportget_model,decodesource_tokens=['i need more power'.split(' '),'eat jujube and pill'.split(' '),]target_tokens=[list('我要更多的抛瓦'),list('吃枣💊'),]# Generate dictionariesdefbuild_token_dict(token_list):token_dict={'<PAD>':0,'<START>':1,'<END>':2,}fortokensintoken_list:fortokenintokens:iftokennotintoken_dict:token_dict[token]=len(token_dict)returntoken_dictsource_token_dict=build_token_dict(source_tokens)target_token_dict=build_token_dict(target_tokens)target_token_dict_inv={v:kfork,vintarget_token_dict.items()}# Add special tokensencode_tokens=[['<START>']+tokens+['<END>']fortokensinsource_tokens]decode_tokens=[['<START>']+tokens+['<END>']fortokensintarget_tokens]output_tokens=[tokens+['<END>','<PAD>']fortokensintarget_tokens]# Paddingsource_max_len=max(map(len,encode_tokens))target_max_len=max(map(len,decode_tokens))encode_tokens=[tokens+['<PAD>']*(source_max_len-len(tokens))fortokensinencode_tokens]decode_tokens=[tokens+['<PAD>']*(target_max_len-len(tokens))fortokensindecode_tokens]output_tokens=[tokens+['<PAD>']*(target_max_len-len(tokens))fortokensinoutput_tokens]encode_input=[list(map(lambdax:source_token_dict[x],tokens))fortokensinencode_tokens]decode_input=[list(map(lambdax:target_token_dict[x],tokens))fortokensindecode_tokens]decode_output=[list(map(lambdax:[target_token_dict[x]],tokens))fortokensinoutput_tokens]# Build & fit modelmodel=get_model(token_num=max(len(source_token_dict),len(target_token_dict)),embed_dim=32,encoder_num=2,decoder_num=2,head_num=4,hidden_dim=128,dropout_rate=0.05,use_same_embed=False,# Use different embeddings for different languages)model.compile('adam','sparse_categorical_crossentropy')model.summary()model.fit(x=[np.array(encode_input*1024),np.array(decode_input*1024)],y=np.array(decode_output*1024),epochs=10,batch_size=32,)# Predictdecoded=decode(model,encode_input,start_token=target_token_dict['<START>'],end_token=target_token_dict['<END>'],pad_token=target_token_dict['<PAD>'],)print(''.join(map(lambdax:target_token_dict_inv[x],decoded[0][1:-1])))print(''.join(map(lambdax:target_token_dict_inv[x],decoded[1][1:-1])))DecodeIndecode, the word with top probability is selected as the predicted token by default. You can add randomness by settingtop_kandtemperature:decoded=decode(model,encode_input,start_token=target_token_dict['<START>'],end_token=target_token_dict['<END>'],pad_token=target_token_dict['<PAD>'],top_k=10,temperature=1.0,)print(''.join(map(lambdax:target_token_dict_inv[x],decoded[0][1:-1])))print(''.join(map(lambdax:target_token_dict_inv[x],decoded[1][1:-1])))