转换FastText训练数据格式到Parquet（Polars，KIMI）
User: 你知道fasttext用于训练的数据格式吗？
Kimi: FastText 的训练数据格式非常简单：
- 每一行代表一个训练样本。
- 对于文本分类任务，每行的开头是标签（label） ，前缀是
__label__（两个下划线），后面跟的是文本内容。 - 标签和文本之间用空格或制表符分隔。
- 可以有多个标签，也可以没有标签（用于无监督训练）。
示例
单标签分类
__label__positive 这部电影太棒了
__label__negative 质量很差，不推荐
多标签分类
__label__sports __label__news 今天的足球比赛结果
无监督训练（词向量）
今天天气很好
我喜欢吃苹果
注意事项
- 标签前缀必须是
__label__，FastText 会识别这个前缀。 - 文本部分不需要分词，FastText 会自动处理。
- 文件编码建议使用 UTF-8。
如果你有具体任务（文本分类、词向量训练等），我可以给你更详细的格式说明。
User: 我需要将这些数据转换为parquet格式，需要考虑几点：
- 文本可能换行，所以不能直接逐行读取。每个数据只有一个标签。
- 数据量可能很大，输出文件批量保存在文件夹中，能够设置批量大小和paruqet格式中的row大小。
- 用polars实现。
- 写成命令行调用的形式，用python xxx.py --xxx xxx执行。
- 不要修改标签和文本内容。
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 argparseLABEL_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 recordsdef 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 += blkcontinuecomplete = 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_filerows_per_rg = args.rows_per_row_groupwritten_total = 0file_idx = 0buffer = []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 的工具读取。