File size: 2,706 Bytes
8ebda9e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
# -*- encoding: utf-8 -*-
'''
Copyright 2022 The International Digital Economy Academy (IDEA). CCNL team. All rights reserved.
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.
@File : GAVAEModel.py
@Time : 2022/11/04 11:35
@Author : Liang Yuxin
@Version : 1.0
@Contact : [email protected]
@License : (C)Copyright 2022-2023, CCNL-IDEA
'''
import torch
from transformers.modeling_utils import PreTrainedModel
from transformers.configuration_utils import PretrainedConfig
from fengshen.models.DAVAE.DAVAEModel import DAVAEModel
from fengshen.models.GAVAE.gans_model import gans_process
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class GAVAEPretrainedModel(PreTrainedModel):
def _init_weights(self, module):
""" Initialize the weights """
pass # to bypass the not implement error
class GAVAEModel(GAVAEPretrainedModel):
config_class = PretrainedConfig
def __init__(self, config:PretrainedConfig) -> None:
super().__init__(config)
self.config =config
config.device = device
self.gan = gans_process(self.config)
self.vae_model = DAVAEModel(self.config)
def train_gan(self,encoder_tokenizer,decoder_tokenizer,input_texts):
self.vae_model.set_tokenizers(encoder_tokenizer,decoder_tokenizer)
n = len(input_texts)
inputs_latents = self.vae_model.latent_code_from_text_batch(input_texts)
well_trained_gan = False
while not well_trained_gan:
self.gan_training(inputs_latents)
latent = torch.tensor(self.gan.gen_test(n))
if not latent.isnan().any():
well_trained_gan = True
def generate(self,n):
latent_z = torch.tensor(self.gan.gen_test(n)).to(device)
text = self.vae_model.text_from_latent_code_batch(latent_z,prompt=None)
return text
def gan_training(self,inputs_latents):
for gt in range(self.config.gan_epoch):
x_train,y_train,x_test,y_test,perm = self.gan.ready_cls(inputs_latents)
# sent_output:latent_z inputs_labels:id of class label
self.gan.cls_train(x_train, y_train)
x2_gen, y_gen, s_gen = self.gan.ready_gen(inputs_latents)
# s_gen:sent_output
self.gan.gen_train(x2_gen, y_gen, s_gen, gt)
|