77 lines
2.0 KiB
Python
77 lines
2.0 KiB
Python
|
import polars as pl
|
|||
|
|
import pandas as pd
|
|||
|
|
from pathlib import Path
|
|||
|
|
from typing import Tuple
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_data(file_path: str) -> pl.DataFrame:
|
|||
|
|
"""使用Polars加载数据集"""
|
|||
|
|
# 加载csv文件，处理编码问题
|
|||
|
|
df = pl.read_csv(
|
|||
|
|
file_path,
|
|||
|
|
encoding="latin-1",
|
|||
|
|
ignore_errors=True,
|
|||
|
|
has_header=True
|
|||
|
|
)
|
|||
|
|
return df
|
|||
|
|
|
|||
|
|
|
|||
|
|
def clean_data(df: pl.DataFrame) -> pl.DataFrame:
|
|||
|
|
"""清洗数据集"""
|
|||
|
|
# 查看数据集基本信息
|
|||
|
|
print("原始数据集形状:", df.shape)
|
|||
|
|
print("原始数据集列名:", df.columns)
|
|||
|
|
|
|||
|
|
# 删除不必要的列（最后三列都是空的）
|
|||
|
|
df = df.drop(df.columns[-3:])
|
|||
|
|
|
|||
|
|
# 重命名列名
|
|||
|
|
df = df.rename({
|
|||
|
|
"v1": "label",
|
|||
|
|
"v2": "message"
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 查看清洗后的数据集
|
|||
|
|
print("清洗后数据集形状:", df.shape)
|
|||
|
|
print("清洗后数据集列名:", df.columns)
|
|||
|
|
print("标签分布:", df["label"].value_counts())
|
|||
|
|
|
|||
|
|
return df
|
|||
|
|
|
|||
|
|
|
|||
|
|
def preprocess_data(df: pl.DataFrame) -> Tuple[pl.DataFrame, pl.Series]:
|
|||
|
|
"""预处理数据，准备用于模型训练"""
|
|||
|
|
# 将标签转换为数值（ham=0, spam=1）
|
|||
|
|
df = df.with_columns(
|
|||
|
|
pl.when(pl.col("label") == "spam").then(1).otherwise(0).alias("label")
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 分离特征和标签
|
|||
|
|
X = df.drop("label")
|
|||
|
|
y = df["label"]
|
|||
|
|
|
|||
|
|
return X, y
|
|||
|
|
|
|||
|
|
|
|||
|
|
def save_data(df: pl.DataFrame, file_path: str) -> None:
|
|||
|
|
"""保存处理后的数据集"""
|
|||
|
|
df.write_csv(file_path, index=False)
|
|||
|
|
print(f"数据集已保存到: {file_path}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
# 测试数据处理流程
|
|||
|
|
file_path = "../spam.csv"
|
|||
|
|
# 检查文件是否存在
|
|||
|
|
import os
|
|||
|
|
if not os.path.exists(file_path):
|
|||
|
|
file_path = "./spam.csv"
|
|||
|
|
df = load_data(file_path)
|
|||
|
|
df_cleaned = clean_data(df)
|
|||
|
|
X, y = preprocess_data(df_cleaned)
|
|||
|
|
|
|||
|
|
print("特征数据形状:", X.shape)
|
|||
|
|
print("标签数据形状:", y.shape)
|
|||
|
|
print("前5行数据:")
|
|||
|
|
print(df_cleaned.head())
|