欢迎您访问 最编程 本站为您分享编程语言代码,编程技术文章!
您现在的位置是: 首页

在 MediaEval 数据集上构建数据集(文本部分 - 使用 PLM BERT)

最编程 2024-03-04 18:58:43
...
import random import numpy as np import pandas as pd import torch from transformers import BertModel,BertTokenizer from tqdm.auto import tqdm from torch.utils.data import Dataset import re """参考Game-On论文""" """util.py""" def set_seed(seed_value=42): random.seed(seed_value) np.random.seed(seed_value) # 用于设置生成随机数的种子 torch.manual_seed(seed_value) torch.cuda.manual_seed_all(seed_value) """util.py""" """文本预处理-textGraph.py""" # 文本DataSet类 def text_preprocessing(text): """ - Remove entity mentions (eg. '@united') - Correct errors (eg. '&amp;' to '&') @param text (str): a string to be processed. @return text (Str): the processed string. """ # Remove '@name' text = re.sub(r'(@.*?)[\s]', ' ', text) # Replace '&amp;' with '&' text = re.sub(r'&amp;', '&', text) # Remove trailing whitespace text = re.sub(r'\s+', ' ', text).strip() # removes links text = re.sub(r'(?P<url>https?://[^\s]+)', r'', text) # remove @usernames text = re.sub(r"\@(\w+)", "", text) # remove # from #tags text = text.replace('#', '') return text class TextDataset(Dataset): def __init__(self,df,tokenizer): # 包含推文的主文件框架 self.df = df.reset_index(drop=True) # 使用的分词器 self.tokenizer = tokenizer def __len__(self): return len(self.df) def __getitem__(self, idx): if torch.is_tensor(idx): idx = idx.tolist() # 帖子的文本内容 text = self.df['tweetText'][idx] # 作为唯一标识符的id ‘tweetId' unique_id = self.df['tweetId'][idx] # 创建一个空的列表来存储输出结果 input_ids = [] attention_mask = [] # 使用tokenizer分词器 encoded_sent = self.tokenizer.encode_plus( text = text_preprocessing(text), # 这里使用的是预处理的句子,而不是直接对原句子使用tokenizer add_special_tokens=True, # 添加[CLS]以及[SEP]等特殊词元 max_length=512, # 最大截断长度 padding='max_length', # padding的最大长度 return_attention_mask=True, # 返回attention_mask truncation=True # ) # 获取编码效果 input_ids = encoded_sent.get('input_ids') # 获取attention_mask结果 attention_mask = encoded_sent.get('attention_mask') # 将列表转换成张量 input_ids = torch.tensor(input_ids) attention_mask =torch.tensor(attention_mask) return {'input_ids':input_ids,'attention_mask':attention_mask,'unique_id':unique_id} def store_data(bert,device,df,dataset,store_dir): lengths = [] bert.eval() for idx in tqdm(range(len(df))): sample = dataset.__getitem__(idx) print('原始sample[input_ids]和sample[attention_mask]的维度:',sample['input_ids'].shape,sample['attention_mask'].shape) # 升维 input_ids,attention_mask = sample['input_ids'].unsqueeze(0),sample['attention_mask'].unsqueeze(0) input_ids = input_ids.to(device) attention_mask = attention_mask.to(device) # 得到唯一标识属性 unique_id = sample['unique_id'] # 计算token的个数 num_tokens = attention_mask.sum().detach().cpu().item() """不生成新的计算图,而是只做权重更新""" with torch.no_grad(): out = bert(input_ids=input_ids,attention_mask=attention_mask) # last_hidden_state.shape是(batch_size,sequence_length,hidden_size) out_tokens = out.last_hidden_state[:,1:num_tokens,:].detach().cpu().squeeze(0).numpy() # token向量 # 保存token级别表示 filename = f'{emed_dir}{unique_id}.npy' try: np.save(filename, out_tokens) print(f"文件{filename}保存成功") except FileNotFoundError: # 文件不存在,创建新文件并保存 np.save(filename, out_tokens) print(f"文件{filename}创建成功并保存成功") lengths.append(num_tokens) ## Save semantic/ whole text representation # 保存语义 也就是整个文本的表示 out_cls = out.last_hidden_state[:,0,:].unsqueeze(0).detach().cpu().squeeze(0).numpy() ## cls vector filename = f'{emed_dir}{unique_id}_full_text.npy' # 尝试保存.npy文件,如果文件不存在则自动创建 try: np.save(filename, out_cls) print(f"文件{filename}保存成功") except FileNotFoundError: # 文件不存在,创建新文件并保存 np.save(filename, out_cls) print(f"文件{filename}创建成功并保存成功") return lengths if __name__=='__main__': device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # 根目录 root_dir = "./dataset/image-verification-corpus-master/image-verification-corpus-master/mediaeval2015/" emed_dir = './Embedding_File' # 文件路径 train_csv_name = "tweetsTrain.csv" test_csv_name = "tweetsTest.csv" # 加载PLM和分词器 tokenizer = BertTokenizer.from_pretrained('./bert/') bert = BertModel.from_pretrained('./bert/', return_dict=True) bert = bert.to(device) # 用于存储每个推文的Embedding store_dir ="Embed_Post/" # 创建训练数据集的Embedding表示 df_train = pd.read_csv(f'{root_dir}{train_csv_name}') df_train = df_train.dropna().reset_index(drop=True) # 训练数据集的编码结果 train_dataset = TextDataset(df_train,tokenizer) lengths = store_data(bert, device, df_train, train_dataset, store_dir) ## Create graph data for testing set # 为测试集创建Embedding表示 df_test = pd.read_csv(f'{root_dir}{test_csv_name}') df_test = df_test.dropna().reset_index(drop=True) test_dataset = TextDataset(df_test, tokenizer) lengths = store_data(bert, device, df_test, test_dataset, store_dir) """文本预处理-textGraph.py"""

推荐阅读