转换FastText训练数据格式到Parquet（Polars，KIMI）

User: 你知道fasttext用于训练的数据格式吗？

Kimi: FastText 的训练数据格式非常简单：

1. 每一行代表一个训练样本。
2. 对于文本分类任务，每行的开头是标签（label） ，前缀是 __label__（两个下划线），后面跟的是文本内容。
3. 标签和文本之间用空格或制表符分隔。
4. 可以有多个标签，也可以没有标签（用于无监督训练）。

示例

单标签分类

__label__positive 这部电影太棒了
__label__negative 质量很差，不推荐

多标签分类

__label__sports __label__news 今天的足球比赛结果

无监督训练（词向量）

今天天气很好
我喜欢吃苹果

注意事项

• 标签前缀必须是 __label__，FastText 会识别这个前缀。
• 文本部分不需要分词，FastText 会自动处理。
• 文件编码建议使用 UTF-8。

如果你有具体任务（文本分类、词向量训练等），我可以给你更详细的格式说明。

---

User: 我需要将这些数据转换为parquet格式，需要考虑几点：

1. 文本可能换行，所以不能直接逐行读取。每个数据只有一个标签。
2. 数据量可能很大，输出文件批量保存在文件夹中，能够设置批量大小和paruqet格式中的row大小。
3. 用polars实现。
4. 写成命令行调用的形式，用python xxx.py --xxx xxx执行。
5. 不要修改标签和文本内容。

---

Kimi: 下面给出一份可直接执行的 Python 脚本，用 Polars 把 FastText 单标签语料（含换行文本）转成 Parquet，支持：

• 自定义输出目录、批次文件大小（rows_per_file）、Parquet row group 大小（rows_per_row_group）
• 命令行调用：python ft2pq.py --input data.txt --output_dir pq_data --rows_per_file 1000000 --rows_per_row_group 100000

脚本不会改动标签与文本内容，仅做 UTF-8 解码与按 __label__ 切分。

#!/usr/bin/env python3
# ft2pq.py
import os
import re
import polars as pl
from pathlib import Path
import argparse

LABEL_PREFIX = "__label__"
LABEL_RE = re.compile(rf"^{re.escape(LABEL_PREFIX)}(\S+)\s*", flags=re.UNICODE)


def parse_fasttext_chunk(chunk_lines):
    """
    将一段连续行还原成 (label, text) 列表。
    每条样本以 __label__xxx 开头，直到下一条 __label__ 或 chunk 结束。
    """
    records = []
    cur_label, cur_text = None, []

    def flush():
        if cur_label is not None:
            records.append((cur_label, "".join(cur_text).rstrip("\n")))

    for line in chunk_lines:
        match = LABEL_RE.match(line)
        if match:  # 新样本开始
            flush()
            cur_label, cur_text = match.group(1), [line[match.end() :]]
        else:
            if cur_label is not None:  # 续行
                cur_text.append("\n" + line)
            # 如果文件开头没有 label，则丢弃
    flush()
    return records


def fasttext_reader(path, chunk_bytes=50 * 1024 * 1024):
    """
    按字节块读取大文件，保证不截断 UTF-8 字符，并按行边界拆分。
    每次 yield 一个行列表（chunk_lines）。
    """
    with open(path, "rb") as f:
        leftover = b""
        while True:
            blk = f.read(chunk_bytes)
            if not blk:
                if leftover:
                    yield leftover.decode("utf-8").splitlines(keepends=True)
                break
            # 找到最后一个完整行
            idx = blk.rfind(b"\n")
            if idx == -1:
                leftover += blk
                continue
            complete = leftover + blk[: idx + 1]
            leftover = blk[idx + 1 :]
            yield complete.decode("utf-8").splitlines(keepends=True)


def convert(args):
    os.makedirs(args.output_dir, exist_ok=True)
    rows_per_file = args.rows_per_file
    rows_per_rg = args.rows_per_row_group
    written_total = 0
    file_idx = 0
    buffer = []

    schema = {"label": pl.Utf8, "text": pl.Utf8}

    for chunk_lines in fasttext_reader(args.input):
        records = parse_fasttext_chunk(chunk_lines)
        buffer.extend(records)

        while len(buffer) >= rows_per_file:
            df = pl.DataFrame(buffer[:rows_per_file], schema=schema)
            out_path = Path(args.output_dir) / f"part{file_idx:05d}.parquet"
            df.write_parquet(out_path, row_group_size=rows_per_rg, compression="snappy")
            written_total += len(df)
            print(f"Wrote {out_path}  rows={len(df)}  total={written_total}")
            buffer = buffer[rows_per_file:]
            file_idx += 1

    # 尾部剩余
    if buffer:
        df = pl.DataFrame(buffer, schema=schema)
        out_path = Path(args.output_dir) / f"part{file_idx:05d}.parquet"
        df.write_parquet(out_path, row_group_size=rows_per_rg, compression="snappy")
        written_total += len(df)
        print(f"Wrote {out_path}  rows={len(df)}  total={written_total}")

    print("All done.")


def main():
    parser = argparse.ArgumentParser(
        description="Convert FastText single-label corpus (with possible newlines inside text) to Parquet via Polars."
    )
    parser.add_argument("--input", required=True, help="Path to FastText .txt file")
    parser.add_argument("--output_dir", required=True, help="Directory to save parquet files")
    parser.add_argument(
        "--rows_per_file",
        type=int,
        default=1_000_000,
        help="Max rows per output parquet file",
    )
    parser.add_argument(
        "--rows_per_row_group",
        type=int,
        default=100_000,
        help="Row group size inside each parquet file",
    )
    args = parser.parse_args()
    convert(args)


if __name__ == "__main__":
    main()

保存为 ft2pq.py 后直接命令行调用即可：

python ft2pq.py \
  --input corpus.txt \
  --output_dir pq_data \
  --rows_per_file 500000 \
  --rows_per_row_group 50000

运行完毕会在 pq_data/ 下得到若干 part00000.parquet、part00001.parquet … 文件，可直接用 Polars 或任何支持 Parquet 的工具读取。