262 lines
7.9 KiB
Python
262 lines
7.9 KiB
Python
from __future__ import annotations
|
||
|
||
import csv
|
||
import re
|
||
from io import StringIO
|
||
from typing import Any
|
||
|
||
|
||
def _rm_prefix(text: str) -> str:
|
||
return re.sub(
|
||
r"^(问题|答案|回答|user|assistant|Q|A|Question|Answer|问|答)[\t:： ]+",
|
||
"",
|
||
(text or "").strip(),
|
||
flags=re.IGNORECASE,
|
||
)
|
||
|
||
|
||
def _to_qa_chunk(question: str, answer: str, eng: bool = False) -> str:
|
||
qprefix = "Question: " if eng else "问题："
|
||
aprefix = "Answer: " if eng else "回答："
|
||
return "\t".join([qprefix + _rm_prefix(question), aprefix + _rm_prefix(answer)])
|
||
|
||
|
||
def _guess_delimiter(lines: list[str]) -> str:
|
||
comma = 0
|
||
tab = 0
|
||
for line in lines:
|
||
if len(line.split(",")) == 2:
|
||
comma += 1
|
||
if len(line.split("\t")) == 2:
|
||
tab += 1
|
||
return "\t" if tab >= comma else ","
|
||
|
||
|
||
def _extract_pairs_with_delimiter(lines: list[str], delimiter: str) -> list[tuple[str, str]]:
|
||
pairs: list[tuple[str, str]] = []
|
||
question = ""
|
||
answer = ""
|
||
|
||
for line in lines:
|
||
arr = line.split(delimiter)
|
||
if len(arr) != 2:
|
||
if question:
|
||
answer += "\n" + line
|
||
continue
|
||
|
||
if question and answer:
|
||
pairs.append((question, answer))
|
||
question, answer = arr
|
||
|
||
if question:
|
||
pairs.append((question, answer))
|
||
|
||
return [(q.strip(), a.strip()) for q, a in pairs if q.strip()]
|
||
|
||
|
||
def _extract_pairs_from_csv(lines: list[str], delimiter: str) -> list[tuple[str, str]]:
|
||
pairs: list[tuple[str, str]] = []
|
||
question = ""
|
||
answer = ""
|
||
|
||
reader = csv.reader(lines, delimiter=delimiter)
|
||
for row, raw_line in zip(reader, lines, strict=False):
|
||
if len(row) != 2:
|
||
if question:
|
||
answer += "\n" + raw_line
|
||
continue
|
||
|
||
if question and answer:
|
||
pairs.append((question, answer))
|
||
question, answer = row
|
||
|
||
if question:
|
||
pairs.append((question, answer))
|
||
|
||
return [(q.strip(), a.strip()) for q, a in pairs if q.strip()]
|
||
|
||
|
||
def _parse_markdown_table_row(line: str) -> list[str] | None:
|
||
if "|" not in line:
|
||
return None
|
||
|
||
text = line.strip()
|
||
if not text:
|
||
return None
|
||
|
||
if text.startswith("|"):
|
||
text = text[1:]
|
||
if text.endswith("|"):
|
||
text = text[:-1]
|
||
|
||
cells = [cell.strip() for cell in text.split("|")]
|
||
if not cells:
|
||
return None
|
||
|
||
if all(re.fullmatch(r":?-{3,}:?", c.replace(" ", "")) for c in cells if c):
|
||
return None
|
||
|
||
return cells
|
||
|
||
|
||
def _extract_pairs_from_markdown_tables(markdown_content: str) -> list[tuple[str, str]]:
|
||
pairs: list[tuple[str, str]] = []
|
||
|
||
for line in (markdown_content or "").splitlines():
|
||
cells = _parse_markdown_table_row(line)
|
||
if not cells or len(cells) < 2:
|
||
continue
|
||
|
||
question = cells[0]
|
||
answer = cells[1]
|
||
if question and answer:
|
||
pairs.append((question, answer))
|
||
|
||
return pairs
|
||
|
||
|
||
def _md_question_level(line: str) -> tuple[int, str]:
|
||
match = re.match(r"^#*", line)
|
||
if not match:
|
||
return 0, line
|
||
return len(match.group(0)), line.lstrip("#").lstrip()
|
||
|
||
|
||
def _extract_pairs_from_markdown_headings(markdown_content: str) -> list[tuple[str, str]]:
|
||
lines = (markdown_content or "").splitlines()
|
||
if not lines:
|
||
return []
|
||
|
||
pairs: list[tuple[str, str]] = []
|
||
last_answer = ""
|
||
question_stack: list[str] = []
|
||
level_stack: list[int] = []
|
||
code_block = False
|
||
|
||
for line in lines:
|
||
if line.strip().startswith("```"):
|
||
code_block = not code_block
|
||
|
||
question_level = 0
|
||
question = ""
|
||
if not code_block:
|
||
question_level, question = _md_question_level(line)
|
||
|
||
if not question_level or question_level > 6:
|
||
last_answer = f"{last_answer}\n{line}"
|
||
continue
|
||
|
||
if last_answer.strip():
|
||
sum_question = "\n".join(question_stack)
|
||
if sum_question:
|
||
pairs.append((sum_question, last_answer.strip()))
|
||
last_answer = ""
|
||
|
||
while question_stack and question_level <= level_stack[-1]:
|
||
question_stack.pop()
|
||
level_stack.pop()
|
||
|
||
question_stack.append(question)
|
||
level_stack.append(question_level)
|
||
|
||
if last_answer.strip():
|
||
sum_question = "\n".join(question_stack)
|
||
if sum_question:
|
||
pairs.append((sum_question, last_answer.strip()))
|
||
|
||
return pairs
|
||
|
||
|
||
def _extract_pairs_by_prefix(lines: list[str]) -> list[tuple[str, str]]:
|
||
pairs: list[tuple[str, str]] = []
|
||
question = ""
|
||
answer_lines: list[str] = []
|
||
|
||
for line in lines:
|
||
if re.match(r"^(Q|Question|问|问题)\s*[:：]", line, flags=re.IGNORECASE):
|
||
if question:
|
||
pairs.append((question, "\n".join(answer_lines)))
|
||
question = re.sub(r"^(Q|Question|问|问题)\s*[:：]", "", line, flags=re.IGNORECASE).strip()
|
||
answer_lines = []
|
||
continue
|
||
|
||
if re.match(r"^(A|Answer|答|回答)\s*[:：]", line, flags=re.IGNORECASE):
|
||
answer_lines.append(re.sub(r"^(A|Answer|答|回答)\s*[:：]", "", line, flags=re.IGNORECASE).strip())
|
||
continue
|
||
|
||
if question:
|
||
answer_lines.append(line)
|
||
|
||
if question:
|
||
pairs.append((question, "\n".join(answer_lines)))
|
||
|
||
return [(q.strip(), a.strip()) for q, a in pairs if q.strip() and a.strip()]
|
||
|
||
|
||
def _dedupe_pairs(pairs: list[tuple[str, str]]) -> list[tuple[str, str]]:
|
||
res: list[tuple[str, str]] = []
|
||
seen: set[tuple[str, str]] = set()
|
||
|
||
for question, answer in pairs:
|
||
q = question.strip()
|
||
a = answer.strip()
|
||
if not q or not a:
|
||
continue
|
||
key = (q, a)
|
||
if key in seen:
|
||
continue
|
||
seen.add(key)
|
||
res.append((q, a))
|
||
|
||
return res
|
||
|
||
|
||
def chunk_markdown(filename: str, markdown_content: str, parser_config: dict[str, Any] | None = None) -> list[str]:
|
||
parser_config = parser_config or {}
|
||
eng = str(parser_config.get("language", "Chinese")).lower() == "english"
|
||
|
||
suffix = ""
|
||
if filename and "." in filename:
|
||
suffix = "." + filename.lower().split(".")[-1]
|
||
|
||
lines = [line for line in (markdown_content or "").splitlines() if line.strip()]
|
||
pairs: list[tuple[str, str]] = []
|
||
|
||
if suffix in {".xlsx", ".xls"}:
|
||
pairs.extend(_extract_pairs_from_markdown_tables(markdown_content))
|
||
if not pairs:
|
||
delimiter = _guess_delimiter(lines)
|
||
pairs.extend(_extract_pairs_with_delimiter(lines, delimiter))
|
||
elif suffix == ".csv":
|
||
pairs.extend(_extract_pairs_from_markdown_tables(markdown_content))
|
||
delimiter = "\t" if any("\t" in line for line in lines) else ","
|
||
pairs.extend(_extract_pairs_from_csv(lines, delimiter))
|
||
elif suffix == ".txt":
|
||
delimiter = _guess_delimiter(lines)
|
||
pairs.extend(_extract_pairs_with_delimiter(lines, delimiter))
|
||
elif suffix in {".md", ".markdown", ".mdx"}:
|
||
pairs.extend(_extract_pairs_from_markdown_headings(markdown_content))
|
||
pairs.extend(_extract_pairs_from_markdown_tables(markdown_content))
|
||
elif suffix == ".docx":
|
||
pairs.extend(_extract_pairs_from_markdown_headings(markdown_content))
|
||
pairs.extend(_extract_pairs_from_markdown_tables(markdown_content))
|
||
else:
|
||
pairs.extend(_extract_pairs_from_markdown_headings(markdown_content))
|
||
pairs.extend(_extract_pairs_from_markdown_tables(markdown_content))
|
||
pairs.extend(_extract_pairs_by_prefix(lines))
|
||
if not pairs:
|
||
delimiter = _guess_delimiter(lines)
|
||
pairs.extend(_extract_pairs_with_delimiter(lines, delimiter))
|
||
|
||
pairs = _dedupe_pairs(pairs)
|
||
|
||
if not pairs and lines:
|
||
# 最后兜底：把内容按 2 行一组构成问答
|
||
for i in range(0, len(lines), 2):
|
||
q = lines[i]
|
||
a = lines[i + 1] if i + 1 < len(lines) else ""
|
||
if q.strip() and a.strip():
|
||
pairs.append((q, a))
|
||
|
||
return [_to_qa_chunk(q, a, eng=eng) for q, a in pairs]
|