semantic-compressor 1.3__py3-none-any.whl → 1.4__py3-none-any.whl

compressor/semantic.py

@@ -6,14 +6,24 @@ from onnxruntime_extensions import get_library_path
 from compressor.minbpe.regex import RegexTokenizer
 from nltk.tokenize import sent_tokenize
 from multiprocessing import cpu_count
+from spellchecker import SpellChecker
+from nltk.stem import PorterStemmer
+from nltk.stem import RSLPStemmer
 from collections import Counter
 import onnxruntime as ort
+import nltk
+
+# Inicializando os stemmers
+stemmer_english = PorterStemmer()
+stemmer_portuguese = RSLPStemmer()
 
 tokenizer = RegexTokenizer()
 nltk_data_path = str(importlib.resources.files('compressor').joinpath('resources/nltk_data'))
 
 os.environ['NLTK_DATA'] = nltk_data_path
 
+nltk.download('rslp')
+
 english_stopwords_path = str(importlib.resources.files('compressor').joinpath('resources/en_stopwords.pkl'))
 portuguese_stopwords_path = str(importlib.resources.files('compressor').joinpath('resources/pt_stopwords.pkl'))
 fasttext_model_path = str(importlib.resources.files('compressor').joinpath('resources/lid.176.ftz'))
@@ -230,10 +240,28 @@ def compress_text(text, *, target_token_count=None, compression_rate=0.7, refere
 
     return text
 
+def stem_text(text, lang='en'):
+    if lang == 'en':
+        stems = [stemmer_english.stem(word) for word in text.split()]
+        stemmed_text = " ".join(stems)
+    else:
+        stems = [stemmer_portuguese.stem(word) for word in text.split()]
+        stemmed_text = " ".join(stems)
+
+    return stemmed_text
+
+def correct_spelling(frase, detected_lang="pt"):
+    spell = SpellChecker(language=detected_lang)
+    words = frase.split()
+    fixed = [spell.correction(word) for word in words]
+    return " ".join(fixed)
+
 def find_needle_in_haystack(
         *, haystack: str, needle: str, block_size = 300,
         semantic_embeddings_weight: float = 0.3,
-        textual_embeddings_weight: float = 0.7
+        textual_embeddings_weight: float = 0.7,
+        use_stemming: bool = False,
+        correct_spelling_needle: bool = False
     ):
     """
     Finds the string block in the haystack that contains the needle.
@@ -244,7 +272,9 @@ def find_needle_in_haystack(
         block_size (int, optional): The size of each string block. The needle will be searched in each block. Defaults to 350.
         semantic_embeddings_weight (float, optional): The weight of the semantic embeddings in the similarity calculation. Defaults to 0.3.
         textual_embeddings_weight (float, optional): The weight of the textual embeddings in the similarity calculation. Defaults to 0.7.
-
+        use_stemming (bool, optional): Whether to use stemming for the text. Defaults to False.
+        correct_spelling_needle (bool, optional): Whether to correct the spelling of the needle. Defaults to False.
+        
     Returns:
         str: The string block in the haystack that contains the needle. The size of the needle will be less than or equal to the block size.
     """
@@ -252,14 +282,19 @@ def find_needle_in_haystack(
     try:
         # Split the haystack into blocks
         blocks = structurize_text(haystack, tokens_per_chunk=block_size)
+
+        lang = detect_language(f"{needle}\n\n{haystack}")
+        
+        if correct_spelling_needle:
+            needle = correct_spelling(needle, lang)
         
         # Compute the embeddings of the needle
         needle_semantic_embedding = extract_semantic_embeddings(needle)
-        needle_textual_embedding = extract_textual_embeddings(needle.lower())
+        needle_textual_embedding = extract_textual_embeddings(needle.lower() if not use_stemming else stem_text(needle, lang))
 
         # Compute the embeddings of the haystack (each block)
         haystack_semantic_embeddings = [extract_semantic_embeddings(block) for block in blocks]
-        haystack_textual_embeddings = [extract_textual_embeddings(block.lower()) for block in blocks]
+        haystack_textual_embeddings = [extract_textual_embeddings(block.lower() if not use_stemming else stem_text(block.lower(), lang)) for block in blocks]
 
         # Compute the similarity between the needle and each block
         semantic_similarities = [calculate_similarity(needle_semantic_embedding, block_embedding) for block_embedding in haystack_semantic_embeddings]

{semantic_compressor-1.3.dist-info → semantic_compressor-1.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: semantic_compressor
-Version: 1.3
+Version: 1.4
 Author: Carlo Moro
 Author-email: Carlo Moro <cnmoro@gmail.com>
 Classifier: Programming Language :: Python :: 3
@@ -15,6 +15,7 @@ Requires-Dist: scikit-learn
 Requires-Dist: fasttext
 Requires-Dist: onnxruntime
 Requires-Dist: onnxruntime-extensions
+Requires-Dist: pyspellchecker
 
 ```python
 from compressor.semantic import compress_text, find_needle_in_haystack

{semantic_compressor-1.3.dist-info → semantic_compressor-1.4.dist-info}/RECORD

@@ -1,5 +1,5 @@
 compressor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-compressor/semantic.py,sha256=8MQdV-ZmTBMC3sEIQr565hafxS6v5_A_-dNiqb0R5Xg,12379
+compressor/semantic.py,sha256=OxqzVCAnICKD3W_P3SAe4JbJt-PyOs5VVR-go8taZVI,13701
 compressor/minbpe/__init__.py,sha256=wZ1z2QKkncvGgiZDBc91AP5m7-M-MVenPStKbS6xylE,95
 compressor/minbpe/base.py,sha256=tTKag04RRFnc4ppoieBbDV0V6thzi_ZvZTlhOYIoY7Q,6881
 compressor/minbpe/basic.py,sha256=0kD4tU8l2MZegfPaHMfDo5CnaSzb9i1v9tDBy6GwMbg,2883
@@ -8,8 +8,8 @@ compressor/resources/embedding_model.onnx,sha256=uLBbAfCGEJTwR1yyiK0bMDroruLr6W5
 compressor/resources/en_stopwords.pkl,sha256=Q2PyGQnphPUs_jxN9NMSqp2EQjYv4b4oMJY2aMYvbSY,1310
 compressor/resources/lid.176.ftz,sha256=jzRyz-hzintgmejpmcPL-uDc0VaWqsfXc4qAOdtgPoM,938013
 compressor/resources/pt_stopwords.pkl,sha256=-9bJaxJWjeOFxLHLT9D-rI3XTzGC0iLJfMiwBDnkCYI,1716
-semantic_compressor-1.3.dist-info/LICENSE,sha256=DFRihXonZ3qVRaTrzuXNaDI_-h2jyT2SqWqjtTDHfqI,1067
-semantic_compressor-1.3.dist-info/METADATA,sha256=baw_1lughU6R-9nQ_23COy7DP70ZI6H_DlH3YvrYBRU,6148
-semantic_compressor-1.3.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-semantic_compressor-1.3.dist-info/top_level.txt,sha256=qb2SlKrEmMrQDVrhwxu3Wr7U6JupPXtDGrJpIQr8xSc,11
-semantic_compressor-1.3.dist-info/RECORD,,
+semantic_compressor-1.4.dist-info/LICENSE,sha256=DFRihXonZ3qVRaTrzuXNaDI_-h2jyT2SqWqjtTDHfqI,1067
+semantic_compressor-1.4.dist-info/METADATA,sha256=BEKlYCs7nYakGXQzbC_8_Gz-MKSAXzSp01pAD0HjIS0,6178
+semantic_compressor-1.4.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+semantic_compressor-1.4.dist-info/top_level.txt,sha256=qb2SlKrEmMrQDVrhwxu3Wr7U6JupPXtDGrJpIQr8xSc,11
+semantic_compressor-1.4.dist-info/RECORD,,

{semantic_compressor-1.3.dist-info → semantic_compressor-1.4.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (72.1.0)
+Generator: bdist_wheel (0.43.0)
 Root-Is-Purelib: true
 Tag: py3-none-any