# -*- 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 https://fasttext.cc/docs/en/crawl-vectors.html.
: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