main.py

from data_utils import *
from transformers import AdamW, get_linear_schedule_with_warmup, set_seed
from transformers import BartTokenizer, BartConfig, BertConfig, BertTokenizer, T5Config, T5Tokenizer
from model import BartForConditionalGeneration, model_type
from tqdm import tqdm
from argparse import ArgumentParser
import torch
from torch.utils.data import DataLoader
from torch.nn import DataParallel
import os
from collections import defaultdict
from cal_simlarity import KLDivergence, Similarity
import pickle


def getSimilarText(args):
    config = BartConfig.from_pretrained(
    pretrained_model_name_or_path=args.model_path,
        )
    tokenizer = BartTokenizer.from_pretrained(args.model_path, forced_bos_token_id=0)



    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    train_dev_processor = Data_type[args.data_type](args, file_path=args.input_file_path, tokenizer=tokenizer,rate = args.rate)
    datasets = train_dev_processor.get_all_sims()
    
    #三个数据集，每个数据集同样位置对应着一个句子的改写1， 改写2


    model = model_type[args.model_type].from_pretrained(
    args.model_path,
    from_tf=bool(".ckpt" in args.model_path),
            config=config
    ).to(device)
    
    # if args.n_gpu > 0:
    #     model = DataParallel(model)
    model.eval()
    for i_sim, sim_dataset in enumerate(datasets):

        output_dir = args.output_dir
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)


        test_dataloader = DataLoader(sim_dataset, shuffle=False, batch_size=1)
        
        #得到所有的文本，进行生成

        
        num_training_steps = len(test_dataloader)
        
        global_step = 0
        generated_text = []
        sacas_labels = []
        with tqdm(total=num_training_steps) as t:
            
            for batch in test_dataloader:
                t.set_description(f"Generating sim_sentence-{i_sim}")
                batch = train_dev_processor.wrap_batch(batch, device)

                generated_ids = model.generate(batch['input_ids'], max_length = 256)
                text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
                generated_text.append(text)
                sacas_labels.append(batch['sar_label'][0].detach().cpu())
                # print(len(sacas_labels),len(sacas_labels[0]))
                # print(text,type(text),len(text))
                global_step += 1
                t.update(1)
        sim_sentence = pd.DataFrame()
        sim_sentence['id'] = pd.Series(list(range(len(generated_text))))
        sim_sentence['text'] = pd.Series([t[0] for t in generated_text])
        sim_sentence['label'] = pd.Series([int(label.item()) for label in sacas_labels])

        sim_sentence.to_csv(output_dir+'/sim_'+str(i_sim)+'_sentences.csv',encoding='utf-8', index=None)

def getSimilarityOfSentence(args):

    
    config = BertConfig.from_pretrained(
    pretrained_model_name_or_path=args.model_path,
    )
        
    tokenizer = BertTokenizer.from_pretrained(args.model_path, forced_bos_token_id=0)
        

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    train_dev_processor = Data_type[args.data_type](file_path=args.simi_sar_text_path, tokenizer=tokenizer)
    simili_datasets = train_dev_processor.get_all()
    
    #三个数据集，每个数据集同样位置对应着一个句子的改写1， 改写2


    model = model_type[args.model_type].from_pretrained(
    args.model_path,
    from_tf=bool(".ckpt" in args.model_path),
            config=config
    ).to(device)
    
    # if args.n_gpu > 0:
    #     model = DataParallel(model)
    all_representations = []
    all_labels= None

    model.eval()
    for i_sim, sim_dataset in enumerate(simili_datasets):

        output_dir = os.path.join(args.similarity_result_path)


        test_dataloader = DataLoader(sim_dataset, shuffle=False, batch_size=1)
        
        #得到所有的文本，进行生成

        
        num_training_steps = len(test_dataloader) 
        
        global_step = 0
        text_cls = None
        labels = None
        with tqdm(total=num_training_steps) as t:
            for batch in test_dataloader:
                t.set_description(f"Embedding sim_sentence-{i_sim}")
                batch = train_dev_processor.wrap_batch(batch, device)
                sen_repres = model(**batch)
                    
                # print(sen_repres.shape)
                # print(batch['labels'].shape)

                if text_cls is None:
                    text_cls = sen_repres.detach().cpu()
                    labels= batch['labels'].detach().cpu()
                    # break
                else:
                    text_cls = torch.cat((text_cls, sen_repres.detach().cpu()), dim=0)
                    labels = torch.cat((labels, batch['labels'].detach().cpu()), dim=0)
                        
                    
                global_step += 1
                t.update(1)
        all_labels = labels
        all_representations.append(text_cls)

    print(len(all_representations),all_representations[0].shape)

    assert all_representations[0].shape[0] == all_representations[1].shape[0] and \
        all_representations[0].shape[0] == all_representations[2].shape[0] 

    if 'kl' not in args.simla_type:
        sim = Similarity(temp=1)
        if 'all' not in args.data_type:
            ori_simA_score = sim(all_representations[0], all_representations[1])
            ori_simB_score = sim(all_representations[0], all_representations[2])
            simA_simB_score = sim(all_representations[1], all_representations[2])
            
            score_result = pd.DataFrame()
            score_result['id'] = pd.Series(list(range(all_representations[0].shape[0])))

            score_result['ori_simA']=pd.Series(ori_simA_score)
            score_result['ori_simB']=pd.Series(ori_simB_score)
            score_result['simA_simB']=pd.Series(simA_simB_score)
            score_result['labels'] = pd.Series(all_labels.long())

            score_result.to_csv(args.similarity_result_path, encoding='utf-8', index=None)
        else:
            keys = ['ori', 'simA', 'simB', 'O4A', 'O4B']
            score_result = pd.DataFrame()
            score_result['id'] = pd.Series(list(range(all_representations[0].shape[0])))

            with open("/home/wangrui/Sarcasm/bart-finetune/data/texts/repres4text/5repres.pkl",'wb') as f:
                pickle.dump({'data':torch.cat(all_representations, dim=-1),'labels':labels},f)
            
            for i in range(len(all_representations)):
                for j in range(i+1, len(all_representations)):
                    score_result[keys[i]+'_'+keys[j]] = pd.Series(sim(all_representations[i], all_representations[j]))
            score_result['labels'] = pd.Series(all_labels.long())
            
            score_result.to_csv(args.similarity_result_path, encoding='utf-8', index=None)
    else:
        ori_simA_score = KLDivergence(all_representations[0], all_representations[1])
        ori_simB_score = KLDivergence(all_representations[0], all_representations[2])
        simA_simB_score = KLDivergence(all_representations[1], all_representations[2])
        # ori_simA_score = KLDivergence()
        score_result = pd.DataFrame()
        score_result['id'] = pd.Series(list(range(all_representations[0].shape[0])))

        score_result['ori_simA']=pd.Series(ori_simA_score)
        score_result['ori_simB']=pd.Series(ori_simB_score)
        score_result['simA_simB']=pd.Series(simA_simB_score)
        score_result['labels'] = pd.Series(all_labels.long())

        score_result.to_csv(args.similarity_result_path, encoding='utf-8', index=None)




if __name__ == '__main__':
    
    arg_parser = ArgumentParser()

    arg_parser.add_argument('--input_file_path', type=str, default=8)
    
    
    arg_parser.add_argument('--mask_file_path', type=str, default=8)

    arg_parser.add_argument('--learning_rate', type=float, default=1e-5)
    
    
    arg_parser.add_argument('--model_path', type=str, required=True)
    arg_parser.add_argument('--output_dir', type=str, required=False)
    
    arg_parser.add_argument('--best_metric_type', type=str, default='f1_marco')
    arg_parser.add_argument('--seed', type=int, default=42)
 
    
    
    arg_parser.add_argument('--model_type', type=str, default='con_gene')
    arg_parser.add_argument('--data_type', type=str, default='con_gene')
    arg_parser.add_argument('--rate', type=float, default=0.5)

    arg_parser.add_argument('--simla_type', type=str, default='cosine')

    arg_parser.add_argument('--similarity_result_path', type=str, default='f1_marco')
    arg_parser.add_argument('--simi_sar_text_path', type=str, default='f1_marco')
    
    arg_parser.add_argument('--action', type=str, default='Mask-Generation', help="Mask-Generation or Sentence-repre")
    
    
    args = arg_parser.parse_args()
    set_seed(args.seed)
    
    args.n_gpu = torch.cuda.device_count()
    if args.action == 'Mask-Generation':
        getSimilarText(args)
        #生成mask后的文本
    else:
        getSimilarityOfSentence(args)
        # 根据mask后的文本，分别计算出句子的表征