Source code for pythainlp.augment.lm.fasttext

# -*- coding: utf-8 -*-
# SPDX-FileCopyrightText: 2016-2024 PyThaiNLP Project
# SPDX-License-Identifier: Apache-2.0
import itertools
from typing import List, Tuple
from gensim.models.fasttext import FastText as FastText_gensim
from gensim.models.keyedvectors import KeyedVectors
from pythainlp.tokenize import word_tokenize

[docs]class FastTextAug: """ Text Augment from fastText :param str model_path: path of model file """
[docs] def __init__(self, model_path: str): """ :param str model_path: path of model file """ if model_path.endswith(".bin"): self.model = FastText_gensim.load_facebook_vectors(model_path) elif model_path.endswith(".vec"): self.model = KeyedVectors.load_word2vec_format(model_path) else: self.model = FastText_gensim.load(model_path) self.dict_wv = list(self.model.key_to_index.keys())
[docs] def tokenize(self, text: str) -> List[str]: """ Thai text tokenization for fastText :param str text: Thai text :return: list of words :rtype: List[str] """ return word_tokenize(text, engine="icu")
[docs] def modify_sent(self, sent: str, p: float = 0.7) -> List[List[str]]: """ :param str sent: text of sentence :param float p: probability :rtype: List[List[str]] """ list_sent_new = [] for i in sent: if i in self.dict_wv: w = [j for j, v in self.model.most_similar(i) if v >= p] if w == []: list_sent_new.append([i]) else: list_sent_new.append(w) else: list_sent_new.append([i]) return list_sent_new
[docs] def augment( self, sentence: str, n_sent: int = 1, p: float = 0.7 ) -> List[Tuple[str]]: """ Text Augment from fastText You may want to download the Thai model from :param str sentence: Thai sentence :param int n_sent: number of sentences :param float p: probability of word :return: list of synonyms :rtype: List[Tuple[str]] """ self.sentence = self.tokenize(sentence) self.list_synonym = self.modify_sent(self.sentence, p=p) new_sentences = [] for x in list(itertools.product(*self.list_synonym))[0:n_sent]: new_sentences.append(x) return new_sentences