doc-to-vector-dataset-generator

文档转向量数据集生成器

将文档转换为高质量的向量搜索数据集。

流程步骤

1. 提取文本，支持多种格式（PDF、DOCX、HTML）
2. 清理文本（去除噪声、规范化）
3. 策略性分块（按语义边界）
4. 添加元数据（来源、时间戳、分类）
5. 去重（检测近似重复内容）
6. 质量检查（长度、内容验证）
7. 导出 JSONL（每行一个块）

文本提取

# PDF extraction
import pymupdf

def extract_pdf(filepath: str) -> str:
    doc = pymupdf.open(filepath)
    text = ""
    for page in doc:
        text += page.get_text()
    return text

# Markdown extraction
def extract_markdown(filepath: str) -> str:
    with open(filepath) as f:
        return f.read()

文本清理

import re

def clean_text(text: str) -> str:
    # Remove extra whitespace
    text = re.sub(r'\s+', ' ', text)

    # Remove page numbers
    text = re.sub(r'Page \d+', '', text)

    # Remove URLs (optional)
    text = re.sub(r'http\S+', '', text)

    # Normalize unicode
    text = text.encode('utf-8', 'ignore').decode('utf-8')

    return text.strip()

语义分块

def semantic_chunk(text: str, max_chunk_size: int = 1000) -> List[str]:
    """Chunk at semantic boundaries (paragraphs, sentences)"""
    # Split by paragraphs first
    paragraphs = text.split('\n\n')

    chunks = []
    current_chunk = ""

    for para in paragraphs:
        if len(current_chunk) + len(para) <= max_chunk_size:
            current_chunk += para + "\n\n"
        else:
            if current_chunk:
                chunks.append(current_chunk.strip())
            current_chunk = para + "\n\n"

    if current_chunk:
        chunks.append(current_chunk.strip())

    return chunks

元数据提取

def extract_metadata(filepath: str, chunk: str, chunk_idx: int) -> dict:
    return {
        "source": filepath,
        "chunk_id": f"{hash(filepath)}_{chunk_idx}",
        "chunk_index": chunk_idx,
        "char_count": len(chunk),
        "word_count": len(chunk.split()),
        "created_at": datetime.now().isoformat(),

        # Content classification
        "has_code": bool(re.search(r'```|def |class |function', chunk)),
        "has_table": bool(re.search(r'\|.*\|', chunk)),
        "language": detect_language(chunk),
    }

去重

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

def deduplicate_chunks(chunks: List[dict], threshold: float = 0.95) -> List[dict]:
    """Remove near-duplicate chunks"""
    texts = [c["text"] for c in chunks]

    # Compute TF-IDF vectors
    vectorizer = TfidfVectorizer()
    vectors = vectorizer.fit_transform(texts)

    # Compute pairwise similarity
    similarity_matrix = cosine_similarity(vectors)

    # Find duplicates
    to_remove = set()
    for i in range(len(chunks)):
        if i in to_remove:
            continue
        for j in range(i+1, len(chunks)):
            if similarity_matrix[i][j] > threshold:
                to_remove.add(j)

    # Return unique chunks
    return [c for i, c in enumerate(chunks) if i not in to_remove]

质量检查

def quality_check(chunk: dict) -> bool:
    """Validate chunk quality"""
    text = chunk["text"]

    # Min length check
    if len(text) < 50:
        return False

    # Max length check
    if len(text) > 5000:
        return False

    # Content check (not just numbers/symbols)
    alpha_ratio = sum(c.isalpha() for c in text) / len(text)
    if alpha_ratio < 0.5:
        return False

    # Language check (English only)
    if chunk["metadata"]["language"] != "en":
        return False

    return True

JSONL 导出

import json

def export_jsonl(chunks: List[dict], output_path: str):
    """Export chunks as JSONL (one JSON object per line)"""
    with open(output_path, 'w') as f:
        for chunk in chunks:
            f.write(json.dumps(chunk) + '\n')

# Example output format
{
  "text": "Chunk text content here...",
  "metadata": {
    "source": "docs/auth.md",
    "chunk_id": "abc123_0",
    "chunk_index": 0,
    "char_count": 542,
    "word_count": 89,
    "has_code": true
  }
}

完整流程

def process_documents(input_dir: str, output_path: str):
    all_chunks = []

    # Process each document
    for filepath in glob(f"{input_dir}/**/*.md"):
        # Extract and clean
        text = extract_markdown(filepath)
        text = clean_text(text)

        # Chunk
        chunks = semantic_chunk(text)

        # Add metadata
        for i, chunk in enumerate(chunks):
            chunk_obj = {
                "text": chunk,
                "metadata": extract_metadata(filepath, chunk, i)
            }

            # Quality check
            if quality_check(chunk_obj):
                all_chunks.append(chunk_obj)

    # Deduplicate
    unique_chunks = deduplicate_chunks(all_chunks)

    # Export
    export_jsonl(unique_chunks, output_path)

    print(f"Processed {len(unique_chunks)} chunks")

最佳实践

• 按语义边界分块
• 使用丰富的元数据进行筛选
• 积极去重
• 通过质量检查避免垃圾内容
• 使用 JSONL 格式进行流式处理
• 为数据集添加版本号

输出检查清单

• 支持从所有格式中提取文本
• 已实现清理流程
• 已制定语义分块策略
• 已定义元数据架构
• 已实现去重逻辑
• 已完成质量验证检查
• JSONL 导出格式
• 已记录数据集统计信息