184 lines
8.0 KiB
Python
184 lines
8.0 KiB
Python
"""Чтение и валидация корпуса реплик Stardew Valley."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import pandas as pd
|
||
|
||
from .cleaning import clean_text, count_placeholders, is_valid_line
|
||
|
||
# Русские названия колонок в исходном CSV/XLSX.
|
||
COLUMN_RENAMES = {
|
||
"Реплика": "text",
|
||
"Персонаж": "character",
|
||
"Text": "text",
|
||
"Character": "character",
|
||
}
|
||
|
||
DEFAULT_DATA_PATH = Path(__file__).resolve().parents[1] / "data" / "dialogues.csv"
|
||
DEFAULT_STATS_PATH = Path(__file__).resolve().parents[1] / "data" / "stats_per_character.xlsx"
|
||
|
||
|
||
class CorpusError(ValueError):
|
||
"""Ошибка структуры или кодировки корпуса."""
|
||
|
||
|
||
def _read_csv_robust(path: Path) -> pd.DataFrame:
|
||
"""Пробует несколько кодировок и проверяет наличие нужных колонок.
|
||
|
||
Основная кодировка корпуса — cp1251. UTF-8 добавлен на случай, если файл
|
||
будет пересохранён; koi8-r — как редкий запасной вариант.
|
||
"""
|
||
last_error: Exception | None = None
|
||
for encoding in ("cp1251", "utf-8-sig", "utf-8", "koi8-r"):
|
||
try:
|
||
frame = pd.read_csv(path, sep=";", encoding=encoding, dtype=str)
|
||
except (UnicodeDecodeError, UnicodeError) as error:
|
||
last_error = error
|
||
continue
|
||
except pd.errors.ParserError as error:
|
||
last_error = error
|
||
continue
|
||
|
||
columns = set(frame.columns)
|
||
if {"Реплика", "Персонаж"} <= columns or {"text", "character"} <= columns:
|
||
return frame
|
||
|
||
# Файл прочитался, но это не тот формат. Пробуем следующую кодировку.
|
||
last_error = CorpusError(
|
||
f"В файле {path} не найдены колонки 'Реплика'/'Персонаж' "
|
||
f"или 'text'/'character'. Найдены: {list(frame.columns)}"
|
||
)
|
||
raise CorpusError(f"Не удалось прочитать CSV {path}: {last_error}")
|
||
|
||
|
||
def _read_excel(path: Path) -> pd.DataFrame:
|
||
"""Читает первый лист Excel. Второй лист в исходном файле пустой."""
|
||
try:
|
||
return pd.read_excel(path, sheet_name=0, dtype=str)
|
||
except Exception as error: # openpyxl поднимает разные типы исключений
|
||
raise CorpusError(f"Не удалось прочитать Excel {path}: {error}") from error
|
||
|
||
|
||
def read_corpus(path: str | Path | None = None) -> pd.DataFrame:
|
||
"""Читает корпус и возвращает таблицу с колонками ``text`` и ``character``.
|
||
|
||
Дополнительно сохраняются:
|
||
- ``text_raw`` — исходная реплика после приведения к строке;
|
||
- ``line_id`` — сквозной номер реплики;
|
||
- ``n_chars`` — длина нормализованного текста;
|
||
- ``placeholder_count`` — число плейсхолдеров.
|
||
"""
|
||
path = Path(path) if path is not None else DEFAULT_DATA_PATH
|
||
if not path.exists():
|
||
raise FileNotFoundError(f"Файл корпуса не найден: {path}")
|
||
|
||
suffix = path.suffix.lower()
|
||
if suffix == ".csv":
|
||
frame = _read_csv_robust(path)
|
||
elif suffix in {".xlsx", ".xls"}:
|
||
frame = _read_excel(path)
|
||
else:
|
||
raise CorpusError(f"Неподдерживаемый формат файла: {suffix}")
|
||
|
||
frame = frame.rename(columns=COLUMN_RENAMES)
|
||
required = {"text", "character"}
|
||
if not required <= set(frame.columns):
|
||
raise CorpusError(
|
||
f"Ожидались колонки {sorted(required)}, найдены {list(frame.columns)}"
|
||
)
|
||
|
||
frame = frame[["text", "character"]].copy()
|
||
frame = frame.dropna(subset=["text", "character"])
|
||
frame["text"] = frame["text"].astype(str).str.strip()
|
||
frame["character"] = frame["character"].astype(str).str.strip()
|
||
frame = frame[(frame["text"] != "") & (frame["character"] != "")]
|
||
|
||
frame["text_raw"] = frame["text"]
|
||
frame["text"] = frame["text_raw"].map(clean_text)
|
||
frame = frame[frame["text"].map(is_valid_line)].copy()
|
||
|
||
frame = frame.reset_index(drop=True)
|
||
frame.insert(0, "line_id", range(len(frame)))
|
||
frame["n_chars"] = frame["text"].str.len().astype(int)
|
||
frame["n_words"] = frame["text"].str.findall(r"[A-Za-z0-9']+").map(len).astype(int)
|
||
frame["placeholder_count"] = frame["text_raw"].map(count_placeholders).astype(int)
|
||
|
||
if frame.empty:
|
||
raise CorpusError("После нормализации корпус оказался пустым.")
|
||
|
||
return frame
|
||
|
||
|
||
def read_stats(path: str | Path | None = None) -> pd.DataFrame:
|
||
"""Читает таблицу статистики по персонажам, если она есть."""
|
||
path = Path(path) if path is not None else DEFAULT_STATS_PATH
|
||
if not path.exists():
|
||
return pd.DataFrame()
|
||
return pd.read_excel(path, sheet_name=0)
|
||
|
||
|
||
def validate_corpus(frame: pd.DataFrame) -> dict[str, Any]:
|
||
"""Считает основные характеристики корпуса для отчёта и проверок."""
|
||
counts = frame["character"].value_counts().sort_values(ascending=False)
|
||
word_counts = frame["text"].str.findall(r"[A-Za-z0-9']+").map(len)
|
||
char_counts = frame["text"].str.len()
|
||
|
||
return {
|
||
"n_lines": int(len(frame)),
|
||
"n_characters": int(frame["character"].nunique()),
|
||
"n_words": int(word_counts.sum()),
|
||
"mean_words_per_line": float(word_counts.mean()),
|
||
"median_words_per_line": float(word_counts.median()),
|
||
"min_words_per_line": int(word_counts.min()),
|
||
"max_words_per_line": int(word_counts.max()),
|
||
"mean_chars_per_line": float(char_counts.mean()),
|
||
"median_chars_per_line": float(char_counts.median()),
|
||
"class_counts": counts.to_dict(),
|
||
"n_duplicate_texts": int(frame["text"].duplicated().sum()),
|
||
"n_placeholder_lines": int((frame["placeholder_count"] > 0).sum()),
|
||
"n_placeholders": int(frame["placeholder_count"].sum()),
|
||
}
|
||
|
||
|
||
def compare_with_stats(frame: pd.DataFrame, stats: pd.DataFrame) -> pd.DataFrame:
|
||
"""Сравнивает фактические counts с таблицей stats_per_character.xlsx.
|
||
|
||
Возвращает таблицу с колонками ``character``, ``actual``, ``reported``,
|
||
``difference``. Если таблица статистики пуста, возвращается пустая таблица.
|
||
"""
|
||
if stats.empty or "Персонаж" not in stats.columns:
|
||
return pd.DataFrame()
|
||
|
||
reported_column = "реплик" if "реплик" in stats.columns else stats.columns[1]
|
||
reported = (
|
||
stats[["Персонаж", reported_column]]
|
||
.rename(columns={"Персонаж": "character", reported_column: "reported"})
|
||
.set_index("character")["reported"]
|
||
)
|
||
actual = frame["character"].value_counts()
|
||
comparison = pd.DataFrame({"actual": actual, "reported": reported})
|
||
comparison["difference"] = comparison["actual"] - comparison["reported"]
|
||
return comparison.fillna(0).astype({"actual": int, "reported": int, "difference": int})
|
||
|
||
|
||
def corpus_fingerprint(path: str | Path | None = None) -> str:
|
||
"""SHA-256 файла корпуса для журнала воспроизводимости."""
|
||
path = Path(path) if path is not None else DEFAULT_DATA_PATH
|
||
digest = hashlib.sha256()
|
||
with path.open("rb") as stream:
|
||
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
|
||
digest.update(chunk)
|
||
return digest.hexdigest()
|
||
|
||
|
||
def write_csv(frame: pd.DataFrame, path: str | Path) -> Path:
|
||
"""Сохраняет таблицу в UTF-8 с BOM, чтобы её открывал Excel."""
|
||
path = Path(path)
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
frame.to_csv(path, index=False, encoding="utf-8-sig")
|
||
return path
|