【提前过年吧】来对对联吧，基于transformer

import paddleimport paddlenlpimport tqdmimport numpy as npimport stringimport randomimport matplotlib.pyplot as pltfrom functools import partial

data_in_path="/home/aistudio/data/data110057/fixed_couplets_in.txt"   # 上联路径data_out_path="/home/aistudio/data/data110057/fixed_couplets_out.txt" # 下联路径

# 从文件中读取数据def openfile(src):
    with open(src,'r',encoding="utf-8") as source:
        lines=source.readlines()    return lines

data_in=openfile(data_in_path)
data_out=openfile(data_out_path)
all_data_lines=len(data_in) # 统计对联的总数，为划分数据集做准备

def delete_newline_and_space(lista): 
    newlist=[]    for i in range(len(lista)):
        newlist.append(["<start>"]+lista[i].strip().split()+['<end>'])    return newlist

data_in_nospace=delete_newline_and_space(data_in)
data_out_nospace=delete_newline_and_space(data_out)# 展示处理结果print("上联：",data_in_nospace[0])print("下联",data_out_nospace[0])

couplet_maxlen=max([len(i) for i in data_in_nospace]) # 获取对联的最大长度，作为统一的长度标准。couplet_maxlen

def bulid_cropus(data_in,data_out):
    crpous=[]    for i in data_in:
        crpous.extend(i)    for i in data_out:
        crpous.extend(i)    return crpous

def build_dict(corpus,frequency):
    # 首先统计不同词（汉字）的频率，使用字典记录
    word_freq_dict={}    for ch in corpus:        if ch not in word_freq_dict:
            word_freq_dict[ch]=0
        word_freq_dict[ch]+=1
    
    # 根据频率对字典进行排序
    word_freq_dict=sorted(word_freq_dict.items(),key=lambda x:x[1],reverse=True)
    
    
    word2id_dict={}
    id2word_dict={}    
    # 按照频率，从高到低，开始遍历每个单词，并赋予第一无二的 id
    for word,freq in word_freq_dict:        if freq>frequency:
            curr_id=len(word2id_dict)
            word2id_dict[word]=curr_id
            id2word_dict[curr_id]=word        else: 
            # else 部分在 使 单词 指向unk,对于汉字，我们不设置unk,令frequency=0
            word2id_dict[word]=1
    return word2id_dict,id2word_dict

word_frequency=0word2id_dict,id2word_dict=build_dict(bulid_cropus(data_in_nospace,data_out_nospace),word_frequency)

word_size=len(word2id_dict)
id_size=len(id2word_dict)print("汉字个数：",word_size,"\n id个数：",id_size)

# 将token-->id的字典、id-->token的字典存储到文件中（不是必须的）with open("word2id.txt",'w',encoding='utf-8') as w2i:    for k,v in word2id_dict.items():
        w2i.write(str(k)+","+str(v)+'\n')with open("id2word.txt",'w',encoding='utf-8') as w2i:    for k,v in id2word_dict.items():
        w2i.write(str(k)+","+str(v)+'\n')

def getensor(w2i,datalist,maxlength=couplet_maxlen):
    in_tensor=[]    for lista in datalist:
        in_samll_tensor=[]        for li in lista:
            in_samll_tensor.append(w2i[li])        if len(in_samll_tensor)<maxlength:
            in_samll_tensor+=[w2i['<end>']]*(maxlength-len(in_samll_tensor))
        in_tensor.append(in_samll_tensor)    return np.array(in_tensor)

in_tensor=getensor(word2id_dict,data_in_nospace)
out_tensor=getensor(word2id_dict,data_out_nospace)


test_data_lines=int(all_data_lines*0.05)
val_data_lines=int(all_data_lines*0.1)


test_in_tensor=in_tensor[:test_data_lines]
val_in_tensor=in_tensor[test_data_lines:val_data_lines]
train_in_tensor=in_tensor[val_data_lines:]

test_out_tensor=out_tensor[:test_data_lines]
val_out_tensor=out_tensor[test_data_lines:val_data_lines]
train_out_tensor=out_tensor[val_data_lines:]print("训练集数目：",len(train_in_tensor),"测试集数目：",len(test_in_tensor),"验证集数目：",len(val_in_tensor))

# 1.继承paddle.io.Datasetclass Mydataset(paddle.io.Dataset):
    
    # 2. 构造函数，定义数据集大小
    def __init__(self,first,second):
        super(Mydataset,self).__init__()
        self.first=first
        self.second=second        
    # 3. 实现__getitem__方法，定义指定index时如何获取数据，并返回单条数据（训练数据，对应的标签）
    def __getitem__(self,index):
        return self.first[index],self.second[index]    
    # 4. 实现__len__方法，返回数据集总数目
    def __len__(self):
        return self.first.shape[0]

def prepare_input(inputs,padid):
    src=np.array([inputsub[0] for inputsub in inputs])
    trg=np.array([inputsub[1] for inputsub in inputs])
    trg_mask =(trg[:,:-1]!=padid).astype(paddle.get_default_dtype())    return src,trg[:,:-1],trg[:,1:,np.newaxis],trg_mask

def create_data_loader(dataset):
    data_loader=paddle.io.DataLoader(dataset,batch_sampler=None,drop_last=True,batch_size=BATCH_SIZE,collate_fn=partial(prepare_input, padid=padid))    return data_loader

# 封装数据集BATCH_SIZE=128padid=word2id_dict['<end>']

train_tensor=Mydataset(train_in_tensor,train_out_tensor)
val_tensor=Mydataset(val_in_tensor,val_out_tensor)

train_loader=create_data_loader(train_tensor)
val_loader=create_data_loader(val_tensor)

for i,data in enumerate(val_loader):    for d in data:        print(d.shape)    break

# 为方便调试网络，我们提前定义一些参数embed_dim=256 # 词嵌入embedding的维度latent_dim=2048 # feed forward 前馈神经网络的相关参数num_heads=8 # 多头注意力机制的‘头’数

class TransformerEncoder(paddle.nn.Layer):
    def __init__(self, embed_dim, dense_dim, num_heads):
        super(TransformerEncoder, self).__init__()
        self.embed_dim = embed_dim
        self.dense_dim = dense_dim
        self.num_heads = num_heads
        self.attention = paddle.nn.MultiHeadAttention(num_heads=num_heads, embed_dim=embed_dim, dropout =0.1)

        self.dense_proj =paddle.nn.Sequential(
            paddle.nn.Linear(embed_dim, dense_dim), 
            paddle.nn.ReLU(),
            paddle.nn.Linear(dense_dim, embed_dim) )

        self.layernorm_1 = paddle.nn.LayerNorm(embed_dim)
        self.layernorm_2 = paddle.nn.LayerNorm(embed_dim)
        self.supports_masking = True

    def forward(self, inputs, mask=None):
        padding_mask=None
        if mask is not None:
            padding_mask = paddle.cast(mask[:, np.newaxis, np.newaxis, :], dtype="int32")
    
        attention_output = self.attention(query=inputs, value=inputs, key=inputs, attn_mask=padding_mask)

        proj_input = self.layernorm_1(inputs + attention_output)
        proj_output = self.dense_proj(proj_input)        return self.layernorm_2(proj_input + proj_output)# pencoder=TransformerEncoder(embed_dim, latent_dim, num_heads)# print(pencoder)# inputs=paddle.rand([BATCH_SIZE,34,256])# print("inputs.shape:",inputs.shape)# out=pencoder(inputs)# print("out.shape:",out.shape)

class PositionalEmbedding(paddle.nn.Layer):
    def __init__(self, sequence_length, vocab_size, embed_dim):
        super(PositionalEmbedding, self).__init__()
        self.token_embeddings = paddle.nn.Embedding(num_embeddings =vocab_size, embedding_dim =embed_dim)
        self.position_embeddings = paddle.nn.Embedding(num_embeddings =sequence_length, embedding_dim =embed_dim)

        self.sequence_length = sequence_length
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim    def forward(self, inputs):
        length = inputs.shape[-1]
        positions = paddle.arange(start=0, end=length, step=1)
        embedded_tokens = self.token_embeddings(inputs)
        embedded_positions = self.position_embeddings(positions)        return embedded_tokens + embedded_positions    def compute_mask(self, inputs, mask=None):
        return paddle.not_equal(inputs, 0)# ps=PositionalEmbedding(34,word_size,256)# print(ps)# inputs=paddle.randint(0,word_size,[BATCH_SIZE,34])# print("inputs.shape:",inputs.shape)# out=ps(inputs)# print("out.shape:",out.shape)

class TransformerDecoder(paddle.nn.Layer):
    def __init__(self, embed_dim, latent_dim, num_heads):
        super(TransformerDecoder, self).__init__()
        self.embed_dim = embed_dim
        self.latent_dim = latent_dim
        self.num_heads = num_heads
        self.attention_1 = paddle.nn.MultiHeadAttention(num_heads=num_heads, embed_dim=embed_dim)

        self.attention_2 = paddle.nn.MultiHeadAttention(num_heads=num_heads, embed_dim=embed_dim)

        self.dense_proj = paddle.nn.Sequential(
            paddle.nn.Linear(embed_dim, latent_dim), 
            paddle.nn.ReLU(),
            paddle.nn.Linear(latent_dim, embed_dim) )

        self.layernorm_1 = paddle.nn.LayerNorm(embed_dim)
        self.layernorm_2 = paddle.nn.LayerNorm(embed_dim)
        self.layernorm_3 = paddle.nn.LayerNorm(embed_dim)
        self.supports_masking = True

    def forward(self, inputs, encoder_outputs, mask=None):
        causal_mask = self.get_causal_attention_mask(inputs) #[batch_size, equence_length, sequence_length]
        padding_mask=None
        if mask is not None:
            padding_mask = paddle.cast(mask[:, np.newaxis, :], dtype="int32")
            padding_mask = paddle.minimum(padding_mask, causal_mask)

        attention_output_1 = self.attention_1(query=inputs, value=inputs, key=inputs, attn_mask=causal_mask)
        out_1 = self.layernorm_1(inputs + attention_output_1)

        attention_output_2 = self.attention_2(
            query=out_1,
            value=encoder_outputs,
            key=encoder_outputs,
            attn_mask=padding_mask,
        )

        out_2 = self.layernorm_2(out_1 + attention_output_2)

        proj_output = self.dense_proj(out_2)        return self.layernorm_3(out_2 + proj_output)    def get_causal_attention_mask(self, inputs):
        input_shape = inputs.shape
        batch_size, sequence_length = input_shape[0], input_shape[1]
        i = paddle.arange(sequence_length)[:, np.newaxis]
        j = paddle.arange(sequence_length)
        mask = paddle.cast(i >= j, dtype="int32") #[sequence_length, sequence_length]
        mask = paddle.reshape(mask, (1,1, input_shape[1], input_shape[1])) #[1, equence_length, sequence_length]
        mult = paddle.concat(
            [paddle.to_tensor(BATCH_SIZE,dtype='int32'), paddle.to_tensor([1,1, 1], dtype="int32")],
            axis=0,) #[batch_size,1,1]
        return paddle.tile(mask, mult) #[batch_size, equence_length, sequence_length]# decoder=TransformerDecoder(embed_dim, latent_dim, num_heads)# print(decoder)# inputs=paddle.rand([BATCH_SIZE,34,256])# enout=paddle.rand([BATCH_SIZE,34,256])# out=decoder(inputs,enout)# print("out.shape:",out.shape)

class Transformer(paddle.nn.Layer):
    def __init__(self, embed_dim, latent_dim, num_heads,sequence_length, vocab_size):
        super(Transformer, self).__init__()

        self.ps1=PositionalEmbedding(sequence_length, vocab_size, embed_dim)
        self.encoder=TransformerEncoder(embed_dim, latent_dim, num_heads)

        self.ps2=PositionalEmbedding(sequence_length, vocab_size, embed_dim)
        self.decoder=TransformerDecoder(embed_dim, latent_dim, num_heads) 

        self.drop=paddle.nn.Dropout(p=0.5)

        self.lastLinear=paddle.nn.Linear(embed_dim,vocab_size)

        self.softmax=paddle.nn.Softmax()    def forward(self,encoder_inputs,decoder_inputs):

        # 编码器
        encoder_emb=self.ps1(encoder_inputs)
        encoder_outputs=self.encoder(encoder_emb)        # 解码器
        deocder_emb=self.ps2(decoder_inputs)
        decoder_outputs=self.decoder(deocder_emb,encoder_outputs)        # dropout
        out=self.drop(decoder_outputs)        #最后输出
        out=self.lastLinear(out)        return out


trans=Transformer(embed_dim, latent_dim, num_heads,couplet_maxlen, word_size)
paddle.summary(trans,input_size=[(BATCH_SIZE,34),(BATCH_SIZE,34)],dtypes='int32')

class CrossEntropy(paddle.nn.Layer):
    def __init__(self):
        super(CrossEntropy,self).__init__()    def forward(self,pre,real,trg_mask):

        # 返回的数据类型与pre一致，除了axis维度(未指定则为-1)，其他维度也与pre一致
        # logits=pre,[batch_size,sequence_len,word_size],猜测会进行argmax操作，[batch_size,sequence_len,1]
        # 默认的soft_label为False，lable=real,[bacth_size,sequence_len,1]
        cost=paddle.nn.functional.softmax_with_cross_entropy(logits=pre,label=real)        
        # 删除axis=2 shape上为1的维度
        # 返回结果的形状应为 [batch_size,sequence_len]
        cost=paddle.squeeze(cost,axis=[2])        # trg_mask 的形状[batch_size,suqence_len]
        # * 这个星号应该是对应位置相乘，返回结果的形状 [bathc_szie,sequence_len]
        masked_cost=cost*trg_mask        # paddle.mean 对应轴的对应位置求平均
        return paddle.mean(paddle.mean(masked_cost,axis=[0]))

epochs = 10   trans=Transformer(embed_dim, latent_dim, num_heads,couplet_maxlen, word_size)
model=paddle.Model(trans)


model.prepare(optimizer=paddle.optimizer.Adam(learning_rate=0.001,parameters=model.parameters()),
                loss=CrossEntropy(), 
                metrics=paddlenlp.metrics.Perplexity())

model.fit(train_data=train_loader, 
            epochs=epochs,
             eval_data= val_loader,
             save_dir='./savemodel',
             save_freq=5,
             verbose =0,
             log_freq =2000,
             callbacks=[paddle.callbacks.VisualDL('./log')])

def evalute(eng,maxlen=couplet_maxlen):
    
    encoder_input=paddle.unsqueeze(eng,axis=0)
    decoded_sentence = "<start>"

    def get_pre_tensor(w2i,s,maxlen=maxlen):
        x=[padid]*couplet_maxlen
        lista=s.split()        for i in range(len(lista)):
            x[i]=w2i[lista[i]]        return paddle.to_tensor([x],dtype='int32')        
    for i in range(maxlen):
        decoder_input=get_pre_tensor(word2id_dict,decoded_sentence)
        pre=trans(encoder_input,decoder_input)
        sampled_token_index = np.argmax(pre[0, i, :])
        sampled_token = id2word_dict[sampled_token_index]
        decoded_sentence += " " + sampled_token        if sampled_token == "<end>":            break

    return decoded_sentence

def translate():
    with open('result.txt','w+') as re:        #for i in tqdm.tqdm(range(len(test_in_tensor))):
        for i in range(5):    
            result=evalute(paddle.to_tensor(test_in_tensor[i]))
            re.write(result+'\n')
translate()

with open('result.txt','r') as re:
    pre=re.readlines()for i in range(2):    print('上联: ',"".join(l for l in data_in_nospace[i]))    print('真实的下联：',"".join(l for l in data_out_nospace[i]))    print('预测的下联：',pre[i])