FINBERT2_finetune / dapt_data_process.py
Riverise's picture
Upload folder using huggingface_hub
fc9ae4e verified
import pandas as pd
import os
from typing import List
# --- 配置参数 ---
INPUT_CSV_PATH = "/home/hsichen/part_time/BERT_finetune/dataset_pretrain/Experiment_sentences_training_filtered_part1.csv"
# 输出的纯文本文件路径 (用于 DAPT 脚本中的 DOMAIN_TEXT_FILE)
OUTPUT_TXT_PATH = "/home/hsichen/part_time/BERT_finetune/dataset_pretrain/domain_corpus.txt"
# 纯文本文件的编码
ENCODING = 'utf-8'
def prepare_dapt_data(input_csv_path: str, output_txt_path: str, encoding: str):
"""
从 CSV 文件中提取 'sentence' 列,并保存为纯文本文件,每行一个句子。
Args:
input_csv_path: 原始 CSV 文件的路径。
output_txt_path: 目标纯文本文件的路径。
encoding: 文件编码。
"""
print(f"--- 1. 读取数据: {input_csv_path} ---")
try:
# 尝试读取 CSV 文件
# 假设文件不含复杂编码问题,使用默认读取
df = pd.read_csv(input_csv_path)
except FileNotFoundError:
print(f"错误:输入文件未找到在路径: {input_csv_path}")
return
except Exception as e:
print(f"读取 CSV 文件时发生错误: {e}")
return
# --- 2. 数据处理与清洗 ---
# 检查 'sentence' 列是否存在
if 'sentence' not in df.columns:
print("错误:CSV 文件中未找到 'sentence' 列。请检查列名是否正确。")
return
# 提取 'sentence' 列,并去除 NaN 值
sentences: List[str] = df['sentence'].dropna().astype(str).tolist()
if not sentences:
print("警告:'sentence' 列中没有有效数据,无法生成语料库。")
return
# 简单清洗:去除多余的空格或换行符(如果有的话)
sentences = [s.strip() for s in sentences]
print(f"提取到 {len(sentences)} 条有效句子。")
# --- 3. 保存为纯文本文件 ---
print(f"--- 3. 保存至纯文本文件: {output_txt_path} ---")
# 将句子列表写入文件,每行一个句子
try:
with open(output_txt_path, 'w', encoding=encoding) as f:
f.write('\n'.join(sentences))
print(f"数据成功保存!")
except Exception as e:
print(f"写入文件时发生错误: {e}")
# --- 运行主函数 ---
if __name__ == "__main__":
# 请确保您已安装 pandas: pip install pandas
prepare_dapt_data(INPUT_CSV_PATH, OUTPUT_TXT_PATH, ENCODING)