解锁数据潜力：信息抽取、数据增强与UIE的完美融合

class UIE(nn.Module):

    def __init__(self, encoder):
        """
        init func.
        Args:
            encoder (transformers.AutoModel): backbone, 默认使用 ernie 3.0
        
        Reference:
            https://github.com/PaddlePaddle/PaddleNLP/blob/a12481fc3039fb45ea2dfac3ea43365a07fc4921/model_zoo/uie/model.py
        """
        super().__init__()
        self.encoder = encoder
        hidden_size = 768
        self.linear_start = nn.Linear(hidden_size, 1)
        self.linear_end = nn.Linear(hidden_size, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(
        self,
        input_ids: torch.tensor,
        token_type_ids: torch.tensor,
        attention_mask=None,
        pos_ids=None,
    ) -> tuple:
        """
        forward 函数，返回开始/结束概率向量。
        Args:
            input_ids (torch.tensor): (batch, seq_len)
            token_type_ids (torch.tensor): (batch, seq_len)
            attention_mask (torch.tensor): (batch, seq_len)
            pos_ids (torch.tensor): (batch, seq_len)
        Returns:
            tuple:  start_prob -> (batch, seq_len)
                    end_prob -> (batch, seq_len)
        """
        sequence_output = self.encoder(
            input_ids=input_ids,
            token_type_ids=token_type_ids,
            position_ids=pos_ids,
            attention_mask=attention_mask,
        )["last_hidden_state"]
        start_logits = self.linear_start(sequence_output)       # (batch, seq_len, 1)
        start_logits = torch.squeeze(start_logits, -1)          # (batch, seq_len)
        start_prob = self.sigmoid(start_logits)                 # (batch, seq_len)
        end_logits = self.linear_end(sequence_output)           # (batch, seq_len, 1)
        end_logits = torch.squeeze(end_logits, -1)              # (batch, seq_len)
        end_prob = self.sigmoid(end_logits)                     # (batch, seq_len)
        return start_prob, end_prob

criterion = torch.nn.BCELoss()
...

start_prob, end_prob = model(input_ids=batch['input_ids'].to(args.device),
                                        token_type_ids=batch['token_type_ids'].to(args.device),
                                        attention_mask=batch['attention_mask'].to(args.device))
start_ids = batch['start_ids'].to(torch.float32).to(args.device)    # (batch, seq_len)
end_ids = batch['end_ids'].to(torch.float32).to(args.device)        # (batch, seq_len)
loss_start = criterion(start_prob, start_ids)                       # 起止向量loss -> (1,)
loss_end = criterion(end_prob, end_ids)                             # 结束向量loss -> (1,)
loss = (loss_start + loss_end) / 2.0                                # 求平均 -> (1,)
loss.backward()
...

pip install -r ../requirements.txt


torch
transformers==4.22.1
datasets==2.4.0
evaluate==0.2.2
matplotlib==3.6.0
rich==12.5.1
scikit-learn==1.1.2
requests==2.28.1

{"content": "谭孝曾是谭元寿的长子，也是谭派第六代传人", "result_list": [{"text": "谭元寿", "start": 4, "end": 7}], "prompt": "谭孝曾的父亲"}
{"content": "在圣保罗书院中学毕业后，曾钰成又在中学会考及大学入学考试中名列全港前十名", "result_list": [{"text": "曾钰成", "start": 12, "end": 15}], "prompt": "人物"}
{"content": "在圣保罗书院中学毕业后，曾钰成又在中学会考及大学入学考试中名列全港前十名", "result_list": [{"text": "圣保罗书院", "start": 1, "end": 6}], "prompt": "曾钰成的毕业院校"}
...

{"text": "谭孝曾是谭元寿的长子，也是谭派第六代传人", "entities": [{"id": 42517, "label": "人物", "start_offset": 0, "end_offset": 3, "text": "谭孝曾"}, {"id": 42518, "label": "人物", "start_offset": 4, "end_offset": 7, "text": "谭元寿"}], "relations": [{"id": 0, "from_id": 42517, "to_id": 42518, "type": "父亲"}]}
...

python train.py \
    --pretrained_model "uie-base-zh" \
    --save_dir "checkpoints/DuIE" \
    --train_path "data/DuIE/train.txt" \
    --dev_path "data/DuIE/dev.txt" \
    --img_log_dir "logs/" \
    --img_log_name "UIE Base" \
    --batch_size 32 \
    --max_seq_len 256 \
    --learning_rate 5e-5 \
    --num_train_epochs 20 \
    --logging_steps 10 \
    --valid_steps 100 \
    --device cuda:0

...
  0%|          | 0/1 [00:00<?, ?ba/s]
100%|██████████| 1/1 [00:00<00:00,  6.91ba/s]
100%|██████████| 1/1 [00:00<00:00,  6.89ba/s]
global step 10, epoch: 1, loss: 0.00244, speed: 2.08 step/s
global step 20, epoch: 1, loss: 0.00228, speed: 2.17 step/s
global step 30, epoch: 1, loss: 0.00191, speed: 2.17 step/s
global step 40, epoch: 1, loss: 0.00168, speed: 2.14 step/s
global step 50, epoch: 1, loss: 0.00149, speed: 2.11 step/s
global step 60, epoch: 1, loss: 0.00138, speed: 2.15 step/s
global step 70, epoch: 2, loss: 0.00123, speed: 2.29 step/s
global step 80, epoch: 2, loss: 0.00112, speed: 2.12 step/s
global step 90, epoch: 2, loss: 0.00102, speed: 2.15 step/s
global step 100, epoch: 2, loss: 0.00096, speed: 2.15 step/s
Evaluation precision: 0.80851, recall: 0.84444, F1: 0.82609
best F1 performence has been updated: 0.00000 --> 0.82609
...

 if __name__ == "__main__":
    from rich import print
    sentences = [
        '谭孝曾是谭元寿的长子，也是谭派第六代传人。'
    ]
    
    # NER 示例
    for sentence in sentences:
        ner_example(
            model,
            tokenizer,
            device,
            sentence=sentence, 
            schema=['人物']
        )

    # SPO 抽取示例
    for sentence in sentences:
        information_extract_example(
            model,
            tokenizer,
            device,
            sentence=sentence, 
            schema={
                    '人物': ['父亲'],
                }
        )

python inference.py

[+] NER Results: 
{
    '人物': ['谭孝曾', '谭元寿']
}

[+] Information-Extraction Results: 
{
    '谭孝曾': 
            {
                '父亲': ['谭元寿']
            }, 
    '谭元寿': {
                '父亲': []
            }
}

pip install streamlit==1.17.0

streamlit run web_da.py --server.port 8904

...
device = 'cpu'                                                 # 指定设备
generated_dataset_height = 800                                 # 生成样本展示高度                
max_show_num = 500                                             # 生成样本最大保存行数
max_seq_len = 128                                              # 数据集单句最大长度
batch_size = 128                                               # 负例生成时的batch_size

filling_model_path = '这里'                                     # fine-tuned filling model
...

prompt: 王文铭的年龄     answer: 76岁        -> 正确
prompt: 王文铭的去世年龄  answer: 76岁        -> 错误