classes.c31_corpus

  1from inspect import cleandoc
  2from typing import Literal
  3import re
  4import colorama
  5from loguru import logger
  6
  7try:
  8    import spacy
  9except ImportError:
 10    spacy = None
 11
 12from lib import files, helpers, content
 13from .c34_text_cleaner import KTextCleaner
 14
 15DEBUG = False
 16INPUT_THRESHOLD = 80000
 17"""Maximum number of items to include from the list when creating a KCorpus from a list of strings."""
 18
 19
 20class KCorpus:
 21    _NLP_CACHE: dict[str, object] = {}
 22    _MODEL_BY_LANG = {
 23        "en": "en_core_web_md",
 24        "ru": "ru_core_news_md",
 25    }
 26    _VALID_POS = {
 27        "ADJ",
 28        "ADP",
 29        "ADV",
 30        "AUX",
 31        "CCONJ",
 32        "DET",
 33        "INTJ",
 34        "NOUN",
 35        "NUM",
 36        "PART",
 37        "PRON",
 38        "PROPN",
 39        "PUNCT",
 40        "SCONJ",
 41        "SYM",
 42        "VERB",
 43        "X",
 44    }
 45    _POS_ALIASES = {
 46        "adjective": "ADJ",
 47        "adj": "ADJ",
 48        "adposition": "ADP",
 49        "adp": "ADP",
 50        "adverb": "ADV",
 51        "adv": "ADV",
 52        "auxiliary": "AUX",
 53        "aux": "AUX",
 54        "coordinating_conjunction": "CCONJ",
 55        "cconj": "CCONJ",
 56        "determiner": "DET",
 57        "det": "DET",
 58        "interjection": "INTJ",
 59        "intj": "INTJ",
 60        "noun": "NOUN",
 61        "proper_noun": "PROPN",
 62        "propernoun": "PROPN",
 63        "propn": "PROPN",
 64        "numeral": "NUM",
 65        "num": "NUM",
 66        "particle": "PART",
 67        "part": "PART",
 68        "pronoun": "PRON",
 69        "pron": "PRON",
 70        "punctuation": "PUNCT",
 71        "punct": "PUNCT",
 72        "subordinating_conjunction": "SCONJ",
 73        "sconj": "SCONJ",
 74        "symbol": "SYM",
 75        "sym": "SYM",
 76        "verb": "VERB",
 77        "other": "X",
 78        "x": "X",
 79    }
 80    _DEFAULT_BRIDGE_WORDS = [
 81        "and",
 82        "the",
 83        "of",
 84        "to",
 85        "in",
 86        "for",
 87        "with",
 88        "on",
 89        "at",
 90        "by",
 91        "from",
 92        "as",
 93        "that",
 94        "which",
 95        "while",
 96        "after",
 97        "before",
 98        "during",
 99        "within",
100        "without",
101        "between",
102        "across",
103        "through",
104        "about",
105        "around",
106        "under",
107        "over",
108        "into",
109        "against",
110    ]
111
112    def __init__(self, data: str, maxPerBucket: int = 5000, lang: str = "en"):
113        """
114        Args:
115            data: Input data as a string or file path.
116            maxPerBucket (optional): Maximum number of n-grams per bucket.
117            lang: Language code used for spaCy model routing.
118        """
119
120        if files.isFile(data):
121            data = files.readFile(data)
122
123        self.lang = lang
124        self.limit = maxPerBucket
125
126        self.raw: str = data
127        """The original input data (contents of file if file path was given)."""
128        self.pruned: str = self._toPruned(self.raw)
129        """Cleaned version of the input data."""
130        self.sentences: list[str] = self._toSentences(self.pruned)
131        """List of sentences from the pruned data."""
132        self.lines: list[str] = self._toLines(self.pruned)
133        """List of lines from the pruned data."""
134
135        self._nlp = self._getNlp(self.lang)
136        self._doc = self._nlp(self.pruned)
137
138        self.tokens: list[str] = self._toTokens(self._doc)
139        """List of cleaned, filtered tokens from the pruned data."""
140        self.words: list[str] = self._toWords(self.tokens)
141        """List of unique words from the pruned data."""
142
143        self._phraseCandidates: list[str] | None = None
144        self._phraseBuckets: dict[int, list[str]] = {}
145
146    @staticmethod
147    def fromDataList(
148        data: list[str] = "/usr/share/dict/words",
149        threshold: int = INPUT_THRESHOLD,
150        **kwargs,
151    ) -> "KCorpus":
152        """Alternative constructor to create a KCorpus from a list of strings.
153
154            Args:
155                data: List of strings to create the corpus from.
156                threshold: Maximum number of items to include from the list. If the list exceeds this size, it will be randomly sampled down to this limit.
157
158        Useful for cases where the input data is too large to process as a single string."""
159        if files.isFile(data):
160            data = files.readFileLines(data)
161
162        if len(data) > threshold:
163            logger.info(
164                "Sampling KCorpus data: {} items, discarded {}.",
165                threshold,
166                len(data) - threshold,
167            )
168            data = helpers.sampleList(data, threshold)
169
170        data = KTextCleaner().sanitizeForbidden(data, dropStrategy="word")
171        return KCorpus(" ".join(data), **kwargs)
172
173    def __str__(self) -> str:
174        """Returns a human-readable representation of the KCorpus object."""
175        summaryColors = [
176            colorama.Fore.BLUE,
177            colorama.Fore.MAGENTA,
178            colorama.Fore.RED,
179            colorama.Fore.CYAN,
180            colorama.Fore.YELLOW,
181            colorama.Fore.GREEN,
182            colorama.Fore.WHITE,
183        ]
184        ngramSummary = " ".join(
185            [
186                f"{summaryColors[index % len(summaryColors)]}{colorama.Back.BLACK}{len(self._phraseBuckets[n])}{colorama.Style.RESET_ALL} {n}-grams"
187                for index, n in enumerate(sorted(self._phraseBuckets.keys()))
188            ]
189        )
190        if not ngramSummary:
191            ngramSummary = f"{colorama.Fore.WHITE}no n-gram buckets cached"
192
193        return cleandoc(
194            f"""
195            {colorama.Fore.BLACK}{colorama.Back.LIGHTGREEN_EX}KCorpus{colorama.Style.RESET_ALL} {len(self.tokens)} tokens {colorama.Fore.YELLOW}{len(self.words)} words {colorama.Fore.GREEN}{len(self.sentences)} sentences {colorama.Fore.CYAN}{len(self.lines)} lines {ngramSummary}{colorama.Style.RESET_ALL}
196            """
197        )
198
199    # Internals
200    def _toPruned(self, data: str) -> str:
201        """Prune missing glyphs."""
202        # Data already cleaned by now
203        dataPruned = content.omitMissing(input=data, mode="words", debug=0)
204
205        return dataPruned.strip()
206
207    def _toSentences(self, data: str) -> list[str]:
208        """Returns cleaned data split into sentences."""
209        return content.splitStringToSentences(data)
210
211    def _toLines(self, data: str) -> list[str]:
212        """Returns cleaned data split into lines."""
213        lines = data.split("\n")
214        lines = helpers.dedupe(lines)
215        return lines
216
217    def _toTokens(self, doc) -> list[str]:
218        """Returns cleaned, filtered words."""
219
220        wordsAll = [
221            token.text
222            for token in doc
223            if not token.is_space and not token.is_punct and not token.like_num
224        ]
225        wordsFiltered = [token for token in wordsAll if not self._isStopWord(token)]
226        # ? Sanitize with my own defined blacklist
227        wordsSanitized = KTextCleaner().sanitizeForbidden(
228            wordsFiltered, dropStrategy="word"
229        )
230        wordsClean = self._cleanTokens(wordsSanitized)
231
232        return wordsClean
233
234    def _cleanTokens(self, tokens: list[str]) -> list[str]:
235        """Normalize token strings and remove empty fragments."""
236
237        removals = [
238            # "-"
239            r"^\W+$",
240            # "- Hello" => "Hello"
241            # "Hello -" => "Hello"
242            r"^[.,']\s+|\s+[.,']$",
243            # "'s" => None
244            r"^\W+[A-Za-z]$",
245        ]
246
247        def cleanToken(token):
248            for removal in removals:
249                matches = re.compile(rf"{removal}").findall(token)
250                if matches:
251                    before = token
252                    for match in matches:
253                        token = re.sub(rf"{removal}", "", token)
254                        logger.trace(
255                            "[Clean] {} \tfrom\t {} \t=> {}",
256                            match,
257                            before,
258                            token or ("empty"),
259                        )
260
261            # Assume it’s not abbreviation => trim punctuation
262            if len(token) > 4:
263                token = token.strip(".,:;-")
264
265            return token
266
267        tokens = [cleanToken(token) for token in tokens]
268        return [token for token in tokens if token]  # Remove None
269
270    def _normalizePosFilter(self, pos: str | list[str] | None) -> set[str] | None:
271        """Normalize and validate POS filters into spaCy coarse POS labels."""
272        if pos is None:
273            return None
274
275        rawValues = helpers.coerceList(pos)
276        normalizedValues = []
277        for value in rawValues:
278            if not isinstance(value, str):
279                continue
280            key = value.strip().lower()
281            if not key:
282                continue
283
284            normalizedValues.append(self._POS_ALIASES.get(key, key.upper()))
285
286        if not normalizedValues:
287            raise ValueError("POS filter cannot be empty.")
288
289        unknown = sorted(
290            {value for value in normalizedValues if value not in self._VALID_POS}
291        )
292        if unknown:
293            raise ValueError(
294                "Unsupported POS values in KCorpus filter: "
295                f"{unknown}. Supported values: {sorted(self._VALID_POS)}"
296            )
297
298        return set(normalizedValues)
299
300    @classmethod
301    def _getNlp(cls, lang: str):
302        if spacy is None:
303            raise RuntimeError(
304                "spaCy is required for KCorpus. Install with: `pip install spacy`."
305            )
306
307        modelName = cls._MODEL_BY_LANG.get(lang)
308        if not modelName:
309            raise ValueError(
310                f"Unsupported KCorpus language '{lang}'. Supported languages: {list(cls._MODEL_BY_LANG.keys())}"
311            )
312
313        cachedNlp = cls._NLP_CACHE.get(modelName)
314        if cachedNlp:
315            return cachedNlp
316
317        try:
318            nlp = spacy.load(modelName)
319        except Exception as e:
320            raise RuntimeError(
321                f"Missing spaCy model '{modelName}' for language '{lang}'. "
322                f"Install it with: `python3 -m spacy download {modelName}`"
323            ) from e
324
325        cls._NLP_CACHE[modelName] = nlp
326        return nlp
327
328    def _isStopWord(self, token: str) -> bool:
329        lexeme = self._nlp.vocab[token]
330        if lexeme.is_stop:
331            return True
332
333        if token.lower() in self._nlp.Defaults.stop_words:
334            return True
335
336        return False
337
338    def _cleanPhrase(self, text: str) -> str:
339        text = re.sub(r"\s+([,.:;!?)\]])", r"\1", text)
340        text = re.sub(r"([([\"])\s+", r"\1", text)
341        text = re.sub(r"(\w)\s+-\s+(\w)", r"\1-\2", text)
342        text = re.sub(r"\s+", " ", text)
343        text = content.dedupeWords(text)
344        return text.strip().rstrip(",.;:")
345
346    def _extractShortSentences(self, minWords=4, maxWords=10) -> list[str]:
347        phrases = []
348
349        for sentence in self._doc.sents:
350            words = [token for token in sentence if token.is_alpha]
351            if minWords <= len(words) <= maxWords:
352                phrase = self._cleanPhrase(sentence.text)
353                if phrase:
354                    phrases.append(phrase)
355
356        return phrases
357
358    def _extractSvoPhrases(self) -> list[str]:
359        phrases = []
360
361        for token in self._doc:
362            if token.pos_ != "VERB":
363                continue
364
365            subject = None
366            obj = None
367
368            for child in token.children:
369                if child.dep_ in ("nsubj", "nsubjpass") and subject is None:
370                    start = min(t.i for t in child.subtree)
371                    end = max(t.i for t in child.subtree) + 1
372                    subject = self._cleanPhrase(self._doc[start:end].text)
373
374                if child.dep_ in ("dobj", "attr", "pobj") and obj is None:
375                    start = min(t.i for t in child.subtree)
376                    end = max(t.i for t in child.subtree) + 1
377                    obj = self._cleanPhrase(self._doc[start:end].text)
378
379            if subject and obj:
380                phrases.append(self._cleanPhrase(f"{subject} {token.text} {obj}"))
381
382        return phrases
383
384    def _extractRootPhrases(self, minWords=2, maxWords=8) -> list[str]:
385        phrases = []
386
387        for sentence in self._doc.sents:
388            roots = [token for token in sentence if token.dep_ == "ROOT"]
389            if not roots:
390                continue
391
392            root = roots[0]
393            phraseTokens = [root]
394            for child in root.children:
395                if child.dep_ in ("nsubj", "nsubjpass", "dobj", "attr", "prep", "aux"):
396                    phraseTokens.extend(list(child.subtree))
397
398            phraseTokens = sorted(set(phraseTokens), key=lambda t: t.i)
399            if not phraseTokens:
400                continue
401
402            start = phraseTokens[0].i
403            end = phraseTokens[-1].i + 1
404            phrase = self._cleanPhrase(self._doc[start:end].text)
405            wordCount = len([token for token in phrase.split(" ") if token])
406            if minWords <= wordCount <= maxWords:
407                phrases.append(phrase)
408
409        return phrases
410
411    def _extractNounChunks(self, minWords=2, maxWords=4) -> list[str]:
412        if not self._doc.has_annotation("DEP"):
413            return []
414
415        phrases = []
416        try:
417            for chunk in self._doc.noun_chunks:
418                phrase = self._cleanPhrase(chunk.text)
419                words = [token for token in phrase.split(" ") if token]
420                if minWords <= len(words) <= maxWords:
421                    phrases.append(phrase)
422        except (NotImplementedError, ValueError, AttributeError):
423            # Not supported for this language/model (e.g. ru), or missing parse data
424            pass
425        return phrases
426
427    def _buildPhraseCandidates(self) -> list[str]:
428        candidatePhrases = []
429        candidatePhrases.extend(self._extractShortSentences(minWords=2, maxWords=8))
430        candidatePhrases.extend(self._extractRootPhrases(minWords=2, maxWords=8))
431        candidatePhrases.extend(self._extractSvoPhrases())
432        candidatePhrases.extend(self._extractNounChunks(minWords=2, maxWords=4))
433
434        candidatePhrases = [phrase for phrase in candidatePhrases if phrase]
435        candidatePhrases = helpers.dedupe(candidatePhrases)
436        return candidatePhrases
437
438    def _buildPhraseBuckets(self, limit: int) -> dict[int, list[str]]:
439        buckets = {}
440        for n in range(2, 8):
441            buckets[n] = self._buildSinglePhraseBucket(n, limit)
442
443        return buckets
444
445    def _splitPhraseWords(self, text: str) -> list[str]:
446        return re.findall(r"[A-Za-zÀ-ÖØ-öø-ÿ0-9]+(?:[-'][A-Za-zÀ-ÖØ-öø-ÿ0-9]+)*", text)
447
448    def _buildSentenceWindows(self, n: int, limit: int, seen: set[str]) -> list[str]:
449        windows = []
450        for sentence in self.sentences:
451            words = self._splitPhraseWords(sentence)
452            if len(words) < n:
453                continue
454
455            for index in range(len(words) - n + 1):
456                phrase = self._cleanPhrase(" ".join(words[index : index + n]))
457                if not phrase or phrase in seen:
458                    continue
459
460                seen.add(phrase)
461                windows.append(phrase)
462                if len(windows) >= limit:
463                    return windows
464
465        return windows
466
467    def _synthesizeLongPhrases(self, n: int, limit: int, seen: set[str]) -> list[str]:
468        bridgeWords = self._DEFAULT_BRIDGE_WORDS
469        sentenceWords = [
470            self._splitPhraseWords(sentence) for sentence in self.sentences if sentence
471        ]
472        sentenceWords = [words for words in sentenceWords if words]
473
474        synthesized = []
475        for index, words in enumerate(sentenceWords):
476            composed = words[:]
477            bridgeIndex = 0
478            nextIndex = index + 1
479
480            while len(composed) < n and nextIndex < len(sentenceWords):
481                # Insert a single bridge word only when stitching sentence fragments.
482                if composed and len(composed) < n:
483                    composed.append(bridgeWords[bridgeIndex % len(bridgeWords)])
484                    bridgeIndex += 1
485
486                remaining = n - len(composed)
487                if remaining <= 0:
488                    break
489
490                composed.extend(sentenceWords[nextIndex][:remaining])
491                nextIndex += 1
492
493            if len(composed) < n:
494                continue
495
496            phrase = self._cleanPhrase(" ".join(composed[:n]))
497            if not phrase or phrase in seen:
498                continue
499
500            seen.add(phrase)
501            synthesized.append(phrase)
502            if len(synthesized) >= limit:
503                break
504
505        return synthesized
506
507    def _buildSinglePhraseBucket(self, n: int, limit: int) -> list[str]:
508        if n < 2:
509            return []
510
511        if self._phraseCandidates is None:
512            self._phraseCandidates = self._buildPhraseCandidates()
513
514        phrases = []
515        seen = set()
516        for phrase in self._phraseCandidates:
517            wordCount = len([token for token in phrase.split(" ") if token])
518            if wordCount == n:
519                if phrase in seen:
520                    continue
521                seen.add(phrase)
522                phrases.append(phrase)
523
524            if len(phrases) >= limit:
525                return phrases[:limit]
526
527        if len(phrases) < limit:
528            sentenceWindows = self._buildSentenceWindows(n, limit - len(phrases), seen)
529            phrases.extend(sentenceWindows)
530
531        if len(phrases) < limit and n > 4:
532            if DEBUG:
533                logger.info(
534                    "Synthesizing long phrases: n={}, needed={}, seen={}",
535                    n,
536                    limit - len(phrases),
537                    len(seen),
538                )
539            synthesized = self._synthesizeLongPhrases(n, limit - len(phrases), seen)
540            phrases.extend(synthesized)
541
542        return phrases[:limit]
543
544    def _getPhraseBucket(self, n: int) -> list[str]:
545        if n in self._phraseBuckets:
546            return self._phraseBuckets[n]
547
548        self._phraseBuckets[n] = self._buildSinglePhraseBucket(n, self.limit)
549
550        return self._phraseBuckets.get(n, [])
551
552    def _toWords(self, tokens: list[str]) -> list[str]:
553        """
554        Returns unique words from the cleaned data.
555
556        Args:
557            tokens: List of tokens to process.
558        """
559        # ? Remove duplicates
560        words = helpers.dedupe(tokens)
561        return words
562
563    def _toNgrams(
564        self,
565        data: str,
566        limit: int,
567        n: int = 4,
568        mode: Literal["quantity", "score"] = "quantity",
569    ):
570        """
571        Output phrases (n words)
572
573        Args:
574            data: The input data as a list of words.
575            limit: Number of n-grams to return (quantity or score).
576            n: The size of the n-gram (2=bigram, 3=trigram, 4=quadgram).
577            mode: 'quantity' for top-N, 'score' for above a PMI score.
578
579        Returns:
580            List of n-gram phrases.
581        """
582        ngrams = list(self._getPhraseBucket(n))
583        return ngrams[:limit]
584
585    def serveSentences(
586        self,
587        mode: Literal["whole", "separate", "connected"] = "separate",
588        limit: int = 100,
589        shuffle=False,
590    ) -> list[str]:
591        """
592        Serve sentences from the corpus.
593
594        Args:
595            mode: Mode of sentence serving (see below).
596            limit: Limit for sentences or words depending on mode.
597            shuffle: Whether to shuffle sentences.
598
599        Modes:
600        - `whole`
601            - whole sentences in logical succession
602            - `limit` number of sentences
603        - `separate`
604            - Chop sentences one by one
605            - `limit` max words in a sentence
606        - `connected`
607            - Connect whole sentences
608            - `limit` max sentences in a block
609
610        Returns:
611            List of sentences or sentence blocks.
612        """
613        sentences = (
614            helpers.shuffleAtRandomSegment(self.sentences)
615            if shuffle
616            else self.sentences
617        )
618
619        if mode == "whole" and limit is not None:
620            # Avoid out of index if limit exceeds available sentences
621            sentences = sentences[: min(limit, len(sentences))]
622        elif mode in ("separate", "connected"):
623            sentences = content.chopList(sentences, limit, mode, shuffle=False)
624
625        return sentences
626
627    def serveLines(self, clamp: int = None) -> list[str]:
628        """
629        Serve shuffled lines from the corpus, optionally clamped.
630
631        Args:
632            clamp: Maximum number of lines to return.
633
634        Returns:
635            List of lines.
636        """
637        lines = content.chopList(self.lines, clamp, mode="separate", shuffle=True)
638        return lines
639
640    def servePhrases(self, n: str | list | tuple = (1, 2, 3, 4)) -> list[str]:
641        """
642        Serve phrases of specified n-gram sizes.
643
644        Args:
645            n: Word counts to include (one or multiple)
646                - 1 = single words
647                - 2 = bigrams
648                - 3 = trigrams
649                - 4 = quadgrams
650                - 5+ = extended phrase buckets
651
652        Returns:
653            List of phrases.
654        """
655        numbers = helpers.coerceList(n)
656        numbers = [
657            int(number)
658            for number in numbers
659            if isinstance(number, (int, float))
660            or (isinstance(number, str) and number.isdigit())
661        ]
662        if any(number < 1 for number in numbers):
663            raise ValueError("KCorpus phrase bucket sizes must be >= 1.")
664
665        ngrams = {1: self.words}
666        for number in numbers:
667            if number <= 1:
668                continue
669            ngrams[number] = self._getPhraseBucket(number)
670
671        phrases = helpers.flatten([ngrams.get(number, []) for number in numbers])
672        phrases = helpers.dedupe(phrases)
673
674        cleaner = KTextCleaner()
675        phrases = [cleaner.removeOrphanedPunctuation(phrase) for phrase in phrases]
676
677        return phrases
678
679    def serveTokens(
680        self,
681        pos: str | list[str] | None = None,
682        lemma=False,
683        unique=True,
684        alterCase: content.TextCase | None = None,
685        includeStopwords=False,
686        minLength=1,
687    ) -> list[str]:
688        """
689        Serve tokens optionally filtered by spaCy POS type.
690
691        Args:
692            pos: One or more coarse POS labels (for example ADJ, NOUN, VERB)
693                or aliases (for example adjective, noun, verb).
694            lemma: Return lemmatized forms instead of surface tokens.
695            unique: Deduplicate resulting tokens.
696            alterCase: Change the case of the output tokens.
697            includeStopwords: Keep stop words in output.
698            minLength: Minimum token length to keep.
699
700        Returns:
701            List of filtered tokens.
702        """
703        posFilter = self._normalizePosFilter(pos)
704        filtered = []
705
706        for token in self._doc:
707            if token.is_space or token.is_punct or token.like_num or not token.is_alpha:
708                continue
709
710            if posFilter and token.pos_ not in posFilter:
711                continue
712
713            tokenText = token.lemma_ if lemma else token.text
714            tokenText = tokenText.strip()
715            if isinstance(alterCase, str):
716                tokenText = content.changeCase(tokenText, alterCase)
717
718            if not tokenText:
719                continue
720
721            if not includeStopwords and self._isStopWord(tokenText):
722                continue
723
724            filtered.append(tokenText)
725
726        filtered = KTextCleaner().sanitizeForbidden(filtered, dropStrategy="word")
727        filtered = self._cleanTokens(filtered)
728        filtered = [token for token in filtered if len(token) >= minLength]
729
730        if unique:
731            filtered = helpers.dedupe(filtered)
732
733        return filtered