diff --git a/main.py b/main.py index 1c6a3ba..304dd53 100644 --- a/main.py +++ b/main.py @@ -19,25 +19,95 @@ def correct_spanish_text(text: str) -> str: - """Certifies words against the Spanish vocabulary and corrects typos/confusions.""" + """ + Certifies words against the Spanish vocabulary and corrects typos/confusions while preserving: + - acronyms + - proper nouns + - mixed-case names + - alphanumeric identifiers + - numbers + - words containing hyphens + + Only ordinary lowercase words are sent to pyspellchecker. + """ + spell = SpellChecker(language='es') + words = text.split() corrected_words = [] for word in words: - # Extract alphanumeric content to check against dictionary - clean_word = re.sub(r'[^\wáéíóúüñÁÉÍÓÚÜÑ]', '', word) - if clean_word and not clean_word.isdigit(): - # If the word is unknown/misspelled, find the most probable valid Spanish substitute - unknowns = spell.unknown([clean_word.lower()]) - if clean_word.lower() in unknowns: - correction = spell.correction(clean_word.lower()) - if correction: - if clean_word.istitle(): - correction = correction.capitalize() - elif clean_word.isupper(): - correction = correction.upper() - word = word.replace(clean_word, correction) + + # Separate punctuation from the actual token. + match = re.match( + r"^([^\wáéíóúüñÁÉÍÓÚÜÑ]*)([\wáéíóúüñÁÉÍÓÚÜÑ-]+)([^\wáéíóúüñÁÉÍÓÚÜÑ]*)$", + word, + ) + + if not match: + corrected_words.append(word) + continue + + prefix, clean_word, suffix = match.groups() + + # --------------------------------------------------------- + # 1. Numbers / numeric identifiers + # --------------------------------------------------------- + if clean_word.isdigit(): + corrected_words.append(word) + continue + + # --------------------------------------------------------- + # 2. Words containing digits + # --------------------------------------------------------- + if any(char.isdigit() for char in clean_word): + corrected_words.append(word) + continue + + # --------------------------------------------------------- + # 3. Acronyms / all-uppercase words + # --------------------------------------------------------- + if clean_word.isupper() and any(char.isalpha() for char in clean_word): + corrected_words.append(word) + continue + + # --------------------------------------------------------- + # 4. Mixed-case words + # --------------------------------------------------------- + if ( + any(char.isupper() for char in clean_word[1:]) + and any(char.islower() for char in clean_word) + ): + corrected_words.append(word) + continue + + # --------------------------------------------------------- + # 5. Hyphenated words + # --------------------------------------------------------- + if "-" in clean_word: + corrected_words.append(word) + continue + + # --------------------------------------------------------- + # 6. Proper nouns + # --------------------------------------------------------- + if clean_word.istitle(): + corrected_words.append(word) + continue + + # --------------------------------------------------------- + # 7. Perform spell checking. + # --------------------------------------------------------- + normalized_word = clean_word.lower() + + unknowns = spell.unknown([normalized_word]) + + if normalized_word in unknowns: + correction = spell.correction(normalized_word) + + if correction: + word = ( prefix + correction + suffix ) + corrected_words.append(word) return " ".join(corrected_words) diff --git a/requirements.txt b/requirements.txt index 4a38c35..f2e514d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,4 +7,5 @@ numpy pandas fpdf2 huggingface_hub -pytest \ No newline at end of file +pytest +pyspellchecker \ No newline at end of file diff --git a/tests/test_spell_checker.py b/tests/test_spell_checker.py new file mode 100644 index 0000000..45ad5c4 --- /dev/null +++ b/tests/test_spell_checker.py @@ -0,0 +1,45 @@ +import pytest + +from main import correct_spanish_text + + +def test_preserves_acronyms(): + text = "La NASA estudia Marte" + result = correct_spanish_text(text) + + assert "NASA" in result + + +def test_preserves_proper_nouns(): + text = "García vive en Madrid" + result = correct_spanish_text(text) + + assert "García" in result + + +def test_preserves_mixed_case_names(): + text = "McDonald trabaja aquí" + result = correct_spanish_text(text) + + assert "McDonald" in result + + +def test_preserves_hyphenated_names(): + text = "García-López llegó ayer" + result = correct_spanish_text(text) + + assert "García-López" in result + + +def test_preserves_alphanumeric_tokens(): + text = "El vuelo A320 llegó" + result = correct_spanish_text(text) + + assert "A320" in result + + +def test_preserves_numbers(): + text = "Tengo 25 años" + result = correct_spanish_text(text) + + assert "25" in result \ No newline at end of file