提取rtf文件中的高亮文本

const fs = require("fs");
const path = require("path");
const cheerio = require("cheerio");

function parseHtmlFile(filePath) {
  const content = fs.readFileSync(filePath, "utf-8");
  const $ = cheerio.load(content);

  let wordIndex = 0;
  let allText = "";
  // 高亮颜色与类型的对应关系
  const colorMap = {
    silver: "人物",
    yellow: "知识点",
    aqua: "朝代",
    red: "文献",
    lime: "事件",
    gray: "官职",
    fuchsia: "地名",
  };
  const resultMap = {};

  $(".MsoNormal")
    .children()
    .each((i, elem) => {
      // 依次遍历子元素
      const text = $(elem).text();
      allText += text;
      if (elem.name === "b") {
        // 找到高亮文本
        const span = $(elem).find("span");
        const style = span.attr("style");
        const match = style.match(/background:(.*?)(;|$)/); // (;|$) 匹配分号或者字符串结尾
        if (match) {
          let color = match[1];
          color = colorMap[color] || color;
          const spanText = span.text();
          if (!resultMap[color]) {
            resultMap[color] = [];
          }
          resultMap[color].push({
            text: spanText,
            start: wordIndex,
            end: wordIndex + spanText.length,
          });
        }
      }
      wordIndex += text.length;
    });

  return {
    content: allText,
    resultMap: resultMap,
  };
}

function parseHtmlFilesInDirectory(dirPath) {
  const files = fs.readdirSync(dirPath);
  const results = [];

  for (const file of files) {
    if (path.extname(file) === ".html") {
      const filePath = path.join(dirPath, file);
      const result = parseHtmlFile(filePath);
      results.push({ file, result });
    }
  }

  return results;
}

const dirPath = "html";
const results = parseHtmlFilesInDirectory(dirPath);

/**
 * 输出示例结果
 * @param {*} results
 */
function getExampleOutput(results) {
  const outputDir = "html-output";
  if (!fs.existsSync(outputDir)) {
    fs.mkdirSync(outputDir);
  }
  // 写入到文件, utf-8 编码格式
  // fs.writeFileSync('result.json', JSON.stringify(results, null, 2), 'utf-8');
  for (const { file, result } of results) {
    fs.writeFileSync(
      path.join(outputDir, file.replace(".html", ".json")),
      JSON.stringify(result, null, 2),
      "utf-8"
    );
  }
}

/**
 * 输出doccano格式结果, jsonl格式
 * {"id": 1, "text": "text", "labels": [[0, 4, "label"]]}
 * @param {*} results
 */
function getDoccanoOutput(results) {
  const doccanoOutputDir = "doccano-output";
  if (!fs.existsSync(doccanoOutputDir)) {
    fs.mkdirSync(doccanoOutputDir);
  }
  const jsonlData = [];
  for (const { file, result } of results) {
    const id = file.replace(".html", "");
    const { content, resultMap } = result;
    const labels = [];
    for (const key in resultMap) {
      const items = resultMap[key];
      for (const item of items) {
        labels.push([item.start, item.end, key]);
      }
    }
    // 按照起始位置排序
    labels.sort((a, b) => a[0] - b[0]);
    jsonlData.push({
      id,
      text: content,
      label: labels,
      Comments: [],
    });
  }
  fs.writeFileSync(
    path.join(doccanoOutputDir, "data.jsonl"),
    jsonlData.map((item) => JSON.stringify(item)).join("\n"),
    "utf-8"
  );
}

// getExampleOutput(results);
getDoccanoOutput(results);