https://github.com/weizhepei/CasRel中utils.py解读
时间: 2023-10-07 18:09:41 浏览: 46
utils.py 是 CasRel 项目中的一个工具类文件,主要包含一些辅助函数和工具函数,其中包括以下几个函数:
1. load_json_file(): 加载 JSON 文件,返回 JSON 格式的数据。
2. save_json_file(): 将 JSON 数据保存到文件中。
3. load_pickle_file(): 加载 pickle 文件,返回 pickle 格式的数据。
4. save_pickle_file(): 将 pickle 数据保存到文件中。
5. read_file_list(): 读取文件列表,返回文件列表中的所有文件路径。
6. get_now_time(): 获取当前时间,返回时间字符串。
7. get_logger(): 获取 logger 对象,用于记录日志信息。
8. set_seed(): 设置随机种子,保证实验的可重复性。
9. find_entity(): 在文本中查找实体,返回实体在文本中的位置和实体类型。
10. find_all_entity(): 在文本中查找所有实体,返回实体列表,每个实体包含实体在文本中的位置和实体类型。
11. find_all_index(): 在文本中查找子串,返回子串在文本中的起始位置和终止位置。
12. get_tok2char_span(): 获取 token 到字符的映射关系,返回 token 到字符的起始位置和终止位置。
这些函数都是 CasRel 项目中常用的辅助函数,用于实现数据加载、日志记录、随机种子设置、实体查找等功能,可以帮助开发者更加方便地进行实验和数据处理。
相关问题
https://github.com/weizhepei/CasRel中run.py解读
`run.py` 是 `CasRel` 项目的入口文件,用于训练和测试模型。以下是 `run.py` 的主要代码解读和功能说明:
### 导入依赖包和模块
首先,`run.py` 导入了所需的依赖包和模块,包括 `torch`、`numpy`、`argparse`、`logging` 等。
```python
import argparse
import logging
import os
import random
import time
import numpy as np
import torch
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
from casrel import CasRel
from dataset import RE_Dataset
from utils import init_logger, load_tokenizer, set_seed, collate_fn
```
### 解析命令行参数
接下来,`run.py` 解析了命令行参数,包括训练数据路径、模型保存路径、预训练模型路径、学习率等参数。
```python
def set_args():
parser = argparse.ArgumentParser()
parser.add_argument("--train_data", default=None, type=str, required=True,
help="The input training data file (a text file).")
parser.add_argument("--dev_data", default=None, type=str, required=True,
help="The input development data file (a text file).")
parser.add_argument("--test_data", default=None, type=str, required=True,
help="The input testing data file (a text file).")
parser.add_argument("--model_path", default=None, type=str, required=True,
help="Path to save, load model")
parser.add_argument("--pretrain_path", default=None, type=str,
help="Path to pre-trained model")
parser.add_argument("--vocab_path", default=None, type=str, required=True,
help="Path to vocabulary")
parser.add_argument("--batch_size", default=32, type=int,
help="Batch size per GPU/CPU for training.")
parser.add_argument("--gradient_accumulation_steps", default=1, type=int,
help="Number of updates steps to accumulate before performing a backward/update pass.")
parser.add_argument("--learning_rate", default=5e-5, type=float,
help="The initial learning rate for Adam.")
parser.add_argument("--num_train_epochs", default=3, type=int,
help="Total number of training epochs to perform.")
parser.add_argument("--max_seq_length", default=256, type=int,
help="The maximum total input sequence length after tokenization. Sequences longer "
"than this will be truncated, sequences shorter will be padded.")
parser.add_argument("--warmup_proportion", default=0.1, type=float,
help="Linear warmup over warmup_steps.")
parser.add_argument("--weight_decay", default=0.01, type=float,
help="Weight decay if we apply some.")
parser.add_argument("--adam_epsilon", default=1e-8, type=float,
help="Epsilon for Adam optimizer.")
parser.add_argument("--max_grad_norm", default=1.0, type=float,
help="Max gradient norm.")
parser.add_argument("--logging_steps", type=int, default=500,
help="Log every X updates steps.")
parser.add_argument("--save_steps", type=int, default=500,
help="Save checkpoint every X updates steps.")
parser.add_argument("--seed", type=int, default=42,
help="random seed for initialization")
parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu",
help="selected device (default: cuda if available)")
args = parser.parse_args()
return args
```
### 加载数据和模型
接下来,`run.py` 加载了训练、验证和测试数据,以及 `CasRel` 模型。
```python
def main():
args = set_args()
init_logger()
set_seed(args)
tokenizer = load_tokenizer(args.vocab_path)
train_dataset = RE_Dataset(args.train_data, tokenizer, args.max_seq_length)
dev_dataset = RE_Dataset(args.dev_data, tokenizer, args.max_seq_length)
test_dataset = RE_Dataset(args.test_data, tokenizer, args.max_seq_length)
train_sampler = RandomSampler(train_dataset)
train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.batch_size,
collate_fn=collate_fn)
dev_sampler = SequentialSampler(dev_dataset)
dev_dataloader = DataLoader(dev_dataset, sampler=dev_sampler, batch_size=args.batch_size,
collate_fn=collate_fn)
test_sampler = SequentialSampler(test_dataset)
test_dataloader = DataLoader(test_dataset, sampler=test_sampler, batch_size=args.batch_size,
collate_fn=collate_fn)
model = CasRel(args)
if args.pretrain_path:
model.load_state_dict(torch.load(args.pretrain_path, map_location="cpu"))
logging.info(f"load pre-trained model from {args.pretrain_path}")
model.to(args.device)
```
### 训练模型
接下来,`run.py` 开始训练模型,包括前向传播、反向传播、梯度更新等步骤。
```python
optimizer = torch.optim.Adam([{'params': model.bert.parameters(), 'lr': args.learning_rate},
{'params': model.subject_fc.parameters(), 'lr': args.learning_rate},
{'params': model.object_fc.parameters(), 'lr': args.learning_rate},
{'params': model.predicate_fc.parameters(), 'lr': args.learning_rate},
{'params': model.linear.parameters(), 'lr': args.learning_rate}],
lr=args.learning_rate, eps=args.adam_epsilon, weight_decay=args.weight_decay)
total_steps = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
warmup_steps = int(total_steps * args.warmup_proportion)
scheduler = torch.optim.lr_scheduler.LambdaLR(
optimizer,
lr_lambda=lambda epoch: 1 / (1 + 0.05 * (epoch - 1))
)
global_step = 0
best_f1 = 0
for epoch in range(args.num_train_epochs):
for step, batch in enumerate(train_dataloader):
model.train()
batch = tuple(t.to(args.device) for t in batch)
inputs = {
"input_ids": batch[0],
"attention_mask": batch[1],
"token_type_ids": batch[2],
"subj_pos": batch[3],
"obj_pos": batch[4],
"subj_type": batch[5],
"obj_type": batch[6],
"subj_label": batch[7],
"obj_label": batch[8],
"predicate_label": batch[9],
}
outputs = model(**inputs)
loss = outputs[0]
loss.backward()
if (step + 1) % args.gradient_accumulation_steps == 0:
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
optimizer.step()
scheduler.step()
optimizer.zero_grad()
global_step += 1
if global_step % args.logging_steps == 0:
logging.info(f"Epoch:[{epoch + 1}]/[{args.num_train_epochs}] Step:[{global_step}] "
f"Train loss:{loss.item():.6f}")
if global_step % args.save_steps == 0:
f1 = evaluate(model, dev_dataloader, args)
if f1 > best_f1:
best_f1 = f1
torch.save(model.state_dict(), os.path.join(args.model_path, "best_model.bin"))
logging.info(f"Save model at step [{global_step}] with best f1 {best_f1:.4f}")
```
### 测试模型
最后,`run.py` 对模型进行测试,输出模型在测试集上的预测结果。
```python
model.load_state_dict(torch.load(os.path.join(args.model_path, "best_model.bin"), map_location="cpu"))
logging.info(f"load best model from {os.path.join(args.model_path, 'best_model.bin')}")
f1, precision, recall = evaluate(model, test_dataloader, args)
logging.info(f"Test f1:{f1:.4f} precision:{precision:.4f} recall:{recall:.4f}")
```
以上就是 `run.py` 的主要代码解读和功能说明。
https://github.com/weizhepei/CasRel 中model.py的解读
CasRel 是一个用于关系抽取的模型,model.py 中实现了模型的定义和训练过程。下面是对 model.py 中的主要代码进行解读:
1. 定义了一个名为 CasRel 的类,继承自 torch.nn.Module,这个类是整个模型的主体。
2. 初始化函数 \_\_init\_\_ 中定义了模型的各个组成部分,包括:
- word_embedding:使用预训练的词向量作为输入的词嵌入层,输入为词汇表的大小和词向量的维度。
- pos1_embedding 和 pos2_embedding:分别表示实体头和实体尾在句子中的位置,它们也是使用预训练的位置向量作为输入的位置嵌入层,输入为句子长度和位置向量的维度。
- encoder:编码器,使用 BiLSTM 对输入的词向量和位置向量进行编码,输出为隐藏状态和最后一个时间步的状态。
- attention:注意力层,计算实体对之间的关系,使用了多头注意力机制。
- classifier:分类器,使用全连接层将编码后的状态映射到关系分类的结果。
3. forward 函数实现了模型的前向传播过程,主要流程如下:
- 将输入的句子、实体头和实体尾分别经过词嵌入层、位置嵌入层和 BiLSTM 编码器,得到隐藏状态和最后一个时间步的状态。
- 使用注意力层计算实体对之间的关系,得到关系向量。
- 将关系向量和最后一个时间步的状态拼接并输入到分类器中,得到关系分类的结果。
4. loss 函数计算模型的损失,使用交叉熵损失函数,其中预测结果通过 softmax 函数计算。在模型训练的过程中,通过反向传播更新模型参数。
以上是对 CasRel 模型的 model.py 文件的简要解读,希望能够对你有所帮助。