NLP/lab1/stylometry/tokenization.py

"""Токенизация, разбиение на предложения и частеречная разметка.

По умолчанию используется NLTK: ``word_tokenize``, ``sent_tokenize`` и
``pos_tag``. Если ресурсы NLTK недоступны, включается упрощённый резервный
режим на регулярных выражениях, а POS-признаки отключаются. Это позволяет
запускать аудит и часть признаков даже без скачанных моделей.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Sequence

import pandas as pd

from .lexicons import (
    AUXILIARIES,
    FUNCTION_WORDS,
    MODALS,
    NEGATIONS,
    PERSONAL_PRONOUNS,
    SUBORDINATORS,
    STOPWORD_FALLBACK,
)

DEFAULT_NLTK_DIR = Path(__file__).resolve().parents[1] / "nltk_data"

_WORD_RE = re.compile(r"[A-Za-z]+(?:['’][A-Za-z]+)?|\d+|[^\w\s]")
_SENTENCE_RE = re.compile(r"(?<=[.!?])\s+")

_PUNCT_TAGS = {".", ",", ":", "``", "''", "-LRB-", "-RRB-", "HYPH", "NFP"}
_PTB_TAGS = [
    "CC", "CD", "DT", "EX", "FW", "IN", "JJ", "JJR", "JJS", "LS", "MD",
    "NN", "NNS", "NNP", "NNPS", "PDT", "POS", "PRP", "PRP$", "RB",
    "RBR", "RBS", "RP", "SYM", "TO", "UH", "VB", "VBD", "VBG", "VBN",
    "VBP", "VBZ", "WDT", "WP", "WP$", "WRB",
]
UPOS_TAGS = [
    "ADJ", "ADP", "ADV", "AUX", "CONJ", "DET", "INTJ", "NOUN", "NUM",
    "PRON", "PRT", "PUNCT", "SCONJ", "SYM", "VERB", "X",
]

_AUX_WORDS = AUXILIARIES | MODALS | {"be", "been", "being", "am", "is", "are", "was", "were"}
_NEGATION_WORDS = NEGATIONS | {"not", "n't"}


def _check_nltk_resource(resource: str) -> bool:
    try:
        import nltk

        nltk.data.find(resource)
        return True
    except Exception:
        return False


def ensure_nltk_data(
    nltk_data_dir: str | Path | None = None,
    download: bool = False,
) -> bool:
    """Добавляет локальный каталог NLTK в путь поиска и проверяет ресурсы.

    Возвращает ``True``, если доступны и токенизатор, и POS-теггер.
    """
    import nltk

    directory = Path(nltk_data_dir) if nltk_data_dir is not None else DEFAULT_NLTK_DIR
    if directory.exists():
        path = str(directory.resolve())
        if path not in nltk.data.path:
            nltk.data.path.insert(0, path)

    required = {
        "tokenizer": "tokenizers/punkt_tab",
        "tagger": "taggers/averaged_perceptron_tagger_eng",
    }
    if download:
        for package in (
            "punkt",
            "punkt_tab",
            "averaged_perceptron_tagger",
            "averaged_perceptron_tagger_eng",
            "stopwords",
        ):
            nltk.download(package, download_dir=str(directory), quiet=True)

    return all(_check_nltk_resource(resource) for resource in required.values())


def load_stopwords(nltk_data_dir: str | Path | None = None) -> set[str]:
    """Загружает английские стоп-слова NLTK или использует запасной список."""
    try:
        import nltk

        directory = Path(nltk_data_dir) if nltk_data_dir is not None else DEFAULT_NLTK_DIR
        if directory.exists():
            path = str(directory.resolve())
            if path not in nltk.data.path:
                nltk.data.path.insert(0, path)
        return set(nltk.corpus.stopwords.words("english"))
    except Exception:
        return set(STOPWORD_FALLBACK)


def penn_to_upos(token: str, tag: str) -> str:
    """Переводит тег Penn Treebank в Universal POS.

    Часть решений зависит не только от тега, но и от самого слова: например,
    ``do`` может быть смысловым глаголом или вспомогательным, а ``as`` —
    предлогом или подчинительным союзом.
    """
    lower = token.lower()
    if tag in _PUNCT_TAGS or token in {".", ",", "!", "?", ";", ":", "...", "…"}:
        return "PUNCT"
    if tag == "SYM":
        return "SYM"
    if tag == "CC":
        return "CONJ"
    if tag == "IN":
        return "SCONJ" if lower in SUBORDINATORS else "ADP"
    if tag == "TO":
        return "PRT"
    if tag == "RP" or tag == "POS":
        return "PRT"
    if tag == "MD":
        return "AUX"
    if tag == "EX":
        return "PRON"
    if tag in {"WDT", "WP", "WP$"}:
        return "PRON"
    if tag == "WRB":
        return "ADV"
    if tag == "UH":
        return "INTJ"
    if tag == "CD":
        return "NUM"
    if tag in {"DT", "PDT"}:
        return "DET"
    if tag in {"PRP", "PRP$"}:
        return "PRON"
    if tag.startswith("JJ"):
        return "ADJ"
    if tag.startswith("RB"):
        return "ADV"
    if tag.startswith("NN"):
        return "NOUN"
    if tag.startswith("VB"):
        if lower in _AUX_WORDS:
            return "AUX"
        return "VERB"
    if tag == "FW":
        return "X"
    if tag == "LS":
        return "X"

    # Уточнения по слову для распространённых разговорных форм.
    if lower in _NEGATION_WORDS:
        return "PRT"
    if lower in _AUX_WORDS:
        return "AUX"
    if lower in PERSONAL_PRONOUNS:
        return "PRON"
    return "X"


def collapse_punctuation(tags: Sequence[str]) -> list[str]:
    """Схлопывает подряд идущие знаки препинания в один тег PUNCT."""
    result: list[str] = []
    for tag in tags:
        if tag == "PUNCT" and result and result[-1] == "PUNCT":
            continue
        result.append(tag)
    return result


def is_function_word(token: str, upos: str | None = None) -> bool:
    """Определяет, относится ли слово к закрытому классу."""
    lower = token.lower()
    if lower in FUNCTION_WORDS:
        return True
    if upos in {"ADP", "AUX", "CONJ", "DET", "PRON", "PRT", "SCONJ"}:
        return True
    return False


def is_content_word(token: str, upos: str | None = None, stopwords: set[str] | None = None) -> bool:
    """Определяет, является ли слово содержательным для content-LSA."""
    lower = token.lower()
    if not re.search(r"[A-Za-z]", token):
        return False
    if is_function_word(token, upos):
        return False
    if stopwords and lower in stopwords:
        return False
    return True


@dataclass
class AnalyzedLine:
    """Результат токенизации и разметки одной реплики."""

    tokens: list[str] = field(default_factory=list)
    ptb: list[str] = field(default_factory=list)
    upos: list[str] = field(default_factory=list)
    sentences: list[list[str]] = field(default_factory=list)
    pos_text: str = ""
    fw_text: str = ""
    content_text: str = ""


class Tokenizer:
    """Токенизатор с NLTK и резервным режимом."""

    def __init__(
        self,
        use_nltk: bool = True,
        nltk_data_dir: str | Path | None = None,
        strict: bool = False,
    ) -> None:
        self.use_nltk = bool(use_nltk)
        self.nltk_data_dir = Path(nltk_data_dir) if nltk_data_dir is not None else DEFAULT_NLTK_DIR
        self.strict = strict
        self.stopwords = load_stopwords(self.nltk_data_dir)

        if self.use_nltk:
            available = ensure_nltk_data(self.nltk_data_dir)
            if not available:
                message = (
                    "Ресурсы NLTK не найдены. Скачайте punkt_tab и "
                    "averaged_perceptron_tagger_eng или запустите с --no-nltk."
                )
                if self.strict:
                    raise RuntimeError(message)
                print(f"[tokenization] {message} Включаю резервный режим.")
                self.use_nltk = False

        self._pos_available = self.use_nltk and _check_nltk_resource(
            "taggers/averaged_perceptron_tagger_eng"
        )
        self._word_tokenize = None
        self._sent_tokenize = None
        self._pos_tag = None
        if self.use_nltk:
            from nltk import pos_tag
            from nltk.tokenize import sent_tokenize, word_tokenize

            self._word_tokenize = word_tokenize
            self._sent_tokenize = sent_tokenize
            self._pos_tag = pos_tag

    @property
    def pos_available(self) -> bool:
        return self._pos_available

    def tokenize(self, text: str) -> list[str]:
        if self._word_tokenize is not None:
            try:
                return [token for token in self._word_tokenize(text) if token.strip()]
            except Exception:
                pass
        return _WORD_RE.findall(text)

    def sentence_tokenize(self, text: str) -> list[str]:
        if self._sent_tokenize is not None:
            try:
                sentences = [sentence.strip() for sentence in self._sent_tokenize(text)]
                return [sentence for sentence in sentences if sentence]
            except Exception:
                pass
        return [part.strip() for part in _SENTENCE_RE.split(text) if part.strip()]

    def pos_tag(self, tokens: Sequence[str]) -> list[tuple[str, str]]:
        if self._pos_tag is not None:
            try:
                return list(self._pos_tag(list(tokens)))
            except Exception:
                pass
        return [(token, "X") for token in tokens]

    def analyze(self, text: str) -> AnalyzedLine:
        tokens = self.tokenize(text)
        tagged = self.pos_tag(tokens)
        ptb = [tag for _, tag in tagged]
        upos = [penn_to_upos(token, tag) for token, tag in tagged]

        sentences: list[list[str]] = []
        for sentence in self.sentence_tokenize(text):
            sentence_tokens = self.tokenize(sentence)
            if sentence_tokens:
                sentences.append(sentence_tokens)
        if not sentences and tokens:
            sentences = [tokens]

        pos_sequence = collapse_punctuation(upos)
        pos_text = " ".join(pos_sequence) if pos_sequence else "__EMPTY__"

        fw_tokens = [
            token
            for token, tag in zip(tokens, upos)
            if is_function_word(token, tag)
        ]
        fw_text = " ".join(fw_tokens) if fw_tokens else "__EMPTY__"

        content_tokens = [
            token
            for token, tag in zip(tokens, upos)
            if is_content_word(token, tag, self.stopwords)
        ]
        content_text = " ".join(content_tokens) if content_tokens else "__EMPTY__"

        return AnalyzedLine(
            tokens=tokens,
            ptb=ptb,
            upos=upos,
            sentences=sentences,
            pos_text=pos_text,
            fw_text=fw_text,
            content_text=content_text,
        )

    def analyze_frame(
        self,
        frame: pd.DataFrame,
        text_column: str = "text",
        verbose: bool = False,
    ) -> pd.DataFrame:
        """Добавляет к таблице колонки с токенами, тегами и текстами для LSA."""
        if text_column not in frame.columns:
            raise KeyError(f"Нет колонки {text_column!r}")

        records: list[dict[str, object]] = []
        total = len(frame)
        for position, (_, row) in enumerate(frame.iterrows(), start=1):
            analyzed = self.analyze(str(row[text_column]))
            records.append(
                {
                    "tokens": analyzed.tokens,
                    "ptb": analyzed.ptb,
                    "upos": analyzed.upos,
                    "sentences": analyzed.sentences,
                    "pos_text": analyzed.pos_text,
                    "fw_text": analyzed.fw_text,
                    "content_text": analyzed.content_text,
                }
            )
            if verbose and position % 500 == 0:
                print(f"[tokenization] обработано {position} из {total}")

        additions = pd.DataFrame(records, index=frame.index)
        result = pd.concat([frame.copy(), additions], axis=1)
        return result


def upos_tag_list() -> list[str]:
    """Список универсальных частеречных тегов, используемых в признаках."""
    return list(UPOS_TAGS)


def ptb_tag_list() -> list[str]:
    """Список тегов Penn Treebank, используемых в признаках."""
    return list(_PTB_TAGS)