用实体识别模型提取每一条事实性句子的关键词（实体），并保存到 JSON 文件中

示例代码：

# Generate Keywords
import torch
import os
from tqdm import tqdm
import json
import nltk
import numpy as np

from span_marker import SpanMarkerModel

model = SpanMarkerModel.from_pretrained("tomaarsen/span-marker-mbert-base-multinerd")
model.cuda()
# pipe = pipeline("token-classification", model="jasminejwebb/KeywordIdentifier", aggregation_strategy="simple",device= 'cuda')

# Function to write a dictionary to a JSON file
def wr_dict(filename, dic):
    data = []
    if os.path.isfile(filename):
        with open(filename, 'r') as f:
            data = json.load(f)
    data.append(dic)
    with open(filename, 'w') as f:
        json.dump(data, f)
            
# Function to remove a file
def rm_file(file_path):
    if os.path.exists(file_path):
        os.remove(file_path)


with open('datasource/news_filter_fact_out_same.json', 'r') as file:
    data = json.load(file)
    

save_path = 'datasource/news_filter_entity.json'
rm_file(save_path)
save_list = []

fact_count = sum(len(d['fact_list']) for d in data)
print(f'Before: {fact_count}')

# Run inference
for item in tqdm(data):
    fact_list = item['fact_list']
    item['fact_key'] = []
    for sample in fact_list:
        facts = {}
        sample = sample.strip()
        if not sample.endswith('.'):
            continue    
        facts['fact'] = sample
        tokens = nltk.word_tokenize(sample)
        length = len(tokens)
        if length<18:
            continue
        entity_list = []
        entities = model.predict(sample)
        for e in entities:
            entity_list.append(e['span'])
        if len(entity_list)==0:
            continue
        # stripped_text = []
        # for group in generated_text:
        #     if group['entity_group'] == 'LABEL_1':
        #         stripped_text.append(group['word'].strip())
        # generated_text = generated_text[0]['generated_text']
        # generated_text = generated_text.split(',')
        stripped_text = [text.strip() for text in entity_list]

        print(stripped_text)

        facts['keywords'] = stripped_text

        # RAKE
        # rake_nltk_var.extract_keywords_from_text(sample)
        # keywords = rake_nltk_var.get_ranked_phrases()
        # print(facts['keywords'])
        item['fact_key'].append(facts)
    # print(item['fact_key'])
    save_list.append(item)
    
with open(save_path, "w") as json_file:
    json.dump(save_list, json_file)

fact_count = sum(len(d['fact_key']) for d in data)
print(f'After: {fact_count}')

代码的目的是：用实体识别模型提取每一条事实性句子的关键词（实体），并保存到 JSON 文件中。

代码的结构和运行流程：

---

🔧 一、模型加载

from span_marker import SpanMarkerModel
model = SpanMarkerModel.from_pretrained("tomaarsen/span-marker-mbert-base-multinerd")
model.cuda()

• 使用了 SpanMarker 这个库，它是一种更先进的实体识别方法（NER），相比传统的 BIO 标注性能更高。
• tomaarsen/span-marker-mbert-base-multinerd 是一个已经训练好的多语言实体识别模型。
• .cuda() 表示把模型加载到 GPU 上，加快推理速度。

---

📂 二、数据准备

with open('datasource/news_filter_fact_out_same.json', 'r') as file:
    data = json.load(file)

• 这一步读取已经过滤并处理好的新闻数据，每一项 data[i] 是一条新闻，里面有多个事实性句子（fact）。

{
  "title": "...",
  "fact_list": ["xxx.", "yyy.", ...]
}

---

🧹 三、清理旧文件 + 初始化变量

save_path = 'datasource/news_filter_entity.json'
rm_file(save_path)
save_list = []

• rm_file 删除旧的结果文件（如果存在）。
• save_list 用来保存新生成的数据。

---

📊 四、统计原始句子数量

fact_count = sum(len(d['fact_list']) for d in data)
print(f'Before: {fact_count}')

• 统计一共多少个“事实性句子”准备处理。

---

🔁 五、主循环（核心逻辑）

for item in tqdm(data):

对每一条新闻循环：

1. 拿出其中的事实性句子列表：

fact_list = item['fact_list']
item['fact_key'] = []

2. 对每个句子进行处理：

for sample in fact_list:
    sample = sample.strip()
    if not sample.endswith('.'):
        continue    

• 只处理以英文句号结尾的句子（.endswith('.')）。

    tokens = nltk.word_tokenize(sample)
    if len(tokens) < 18:
        continue

• 用 nltk 对句子进行分词，句子太短（小于 18 个词）就跳过，避免提取没意义的实体。

3. 用模型预测实体：

    entities = model.predict(sample)
    for e in entities:
        entity_list.append(e['span'])

• 用 span-marker 模型提取实体。
• 每个实体是一个 e['span'] 字段（即：这段文本里的实体字符串）。
• 例如句子中可能识别出：['United States', 'Donald Trump', 'White House']

4. 保存实体

    facts = {'fact': sample, 'keywords': stripped_text}
    item['fact_key'].append(facts)

• 保存格式是：

{
  "fact": "原始句子",
  "keywords": ["实体1", "实体2", ...]
}

---

💾 六、最终保存

save_list.append(item)
with open(save_path, "w") as json_file:
    json.dump(save_list, json_file)

• 把每一条新闻及其抽取出的实体（关键词）保存到 news_filter_entity.json 文件中。

---

📈 七、统计处理后数据量

fact_count = sum(len(d['fact_key']) for d in data)
print(f'After: {fact_count}')

• 输出提取成功的句子数量。

---

🔚 总结

步骤	作用
加载模型	用 SpanMarker 加载多语言实体识别模型
读取数据	加载事实性句子列表
循环处理	每条句子判断格式、长度，使用模型提取实体
保存数据	每条句子提取到的实体写入 JSON 文件

---

✅ 输出格式（简化）：

{
  "title": "Some News",
  "fact_list": [...],
  "fact_key": [
    {
      "fact": "Joe Biden visited China on Tuesday.",
      "keywords": ["Joe Biden", "China", "Tuesday"]
    },
    ...
  ]
}