Upload tool

tool.py

@@ -1,8 +1,8 @@
 from smolagents.tools import Tool
 import json
+import difflib
 import string
 import pronouncing
-import difflib
 
 class ParodyWordSuggestionTool(Tool):
     name = "parody_word_suggester"
@@ -11,21 +11,57 @@ class ParodyWordSuggestionTool(Tool):
     inputs = {'target': {'type': 'string', 'description': 'The word you want to find rhyming alternatives for'}, 'word_list_str': {'type': 'string', 'description': 'JSON string of word list (e.g. \'["word1", "word2"]\')'}, 'min_similarity': {'type': 'string', 'description': 'Minimum similarity threshold (0.0-1.0)', 'nullable': True}}
     output_type = "string"
 
-    def _has_vowel(self, phone: str) -> bool:
-        """Check if a phone contains a vowel."""
-        VOWELS = ['A', 'E', 'I', 'O', 'U']
-        for vowel in VOWELS:
-            if vowel in phone:
-                return True
-        return False
+    def __init__(self):
+        self.vowel_groups = [
+            {'AH', 'UH', 'AX'},  # 'luck', 'buck', 'cuck'
+            {'AE', 'EH'},        # 'bat', 'bet'
+            {'IY', 'IH'},        # 'seat', 'sit'
+            {'AO', 'AA'},        # 'caught', 'cot'
+            {'UW', 'UH'},        # 'boot', 'put'
+            {'AY', 'EY'},        # 'bite', 'bait'
+            {'OW', 'AO'},        # 'boat', 'bought'
+            {'AW', 'AO'},        # 'bout', 'bought'
+            {'OY', 'OW'},        # 'boy', 'bow'
+            {'ER', 'AXR'},       # 'bird', 'hurt'
+        ]
+
+
+    def _get_last_syllable(self, phones: list) -> tuple:
+        """Extract the last syllable (vowel + remaining consonants)."""
+        last_vowel_idx = -1
+        last_vowel = None
+    
+        # Find the last vowel
+        for i, phone in enumerate(phones):
+            # Strip stress markers for checking
+            base_phone = phone.rstrip('012')
+            for group in self.vowel_groups:
+                if base_phone in group:
+                    last_vowel_idx = i
+                    last_vowel = base_phone
+                    break
+    
+        if last_vowel_idx == -1:
+            return None, []
+        
+        # Get all consonants after the vowel
+        remaining = phones[last_vowel_idx + 1:]
+    
+        return last_vowel, remaining
 
 
-    def _get_rhyme_part(self, phones: list) -> list:
-        """Get the rhyming part of a word (last vowel onwards)."""
-        for i, phone in enumerate(reversed(phones)):
-            if self._has_vowel(phone):
-                return phones[-(i+1):]
-        return phones
+    def _vowels_match(self, v1: str, v2: str) -> bool:
+        """Check if two vowels are in the same group."""
+        v1 = v1.rstrip('012')
+        v2 = v2.rstrip('012')
+    
+        if v1 == v2:
+            return True
+        
+        for group in self.vowel_groups:
+            if v1 in group and v2 in group:
+                return True
+        return False
 
 
     def forward(self, target: str, word_list_str: str, min_similarity: str = "0.5") -> str:
@@ -35,9 +71,14 @@ class ParodyWordSuggestionTool(Tool):
         import json
         from difflib import SequenceMatcher
     
+        # Initialize variables
         target = target.lower().strip(string.punctuation)
         min_similarity = float(min_similarity)
         suggestions = []
+        word_vowel = None
+        word_end = []
+        target_vowel = None
+        target_end = []
     
         # Parse JSON string to list
         try:
@@ -58,7 +99,7 @@ class ParodyWordSuggestionTool(Tool):
     
         target_phones = target_phones[0]
         target_phone_list = target_phones.split()
-        target_rhyme_part = self._get_rhyme_part(target_phone_list)
+        target_vowel, target_end = self._get_last_syllable(target_phone_list)
     
         # Check each word
         for word in words:
@@ -67,11 +108,19 @@ class ParodyWordSuggestionTool(Tool):
             if phones:
                 word_phones = phones[0]
                 word_phone_list = word_phones.split()
-                word_rhyme_part = self._get_rhyme_part(word_phone_list)
+                word_vowel, word_end = self._get_last_syllable(word_phone_list)
             
                 # 1. Rhyme score (most important - 60%)
-                # Perfect rhyme if the rhyming parts match exactly
-                rhyme_score = 1.0 if word_rhyme_part == target_rhyme_part else 0.0
+                rhyme_score = 0.0
+                if word_vowel and target_vowel:
+                    # Check vowel match
+                    if self._vowels_match(word_vowel, target_vowel):
+                        # Perfect rhyme if endings match too
+                        if word_end == target_end:
+                            rhyme_score = 1.0
+                        # Partial rhyme if just the vowel matches
+                        else:
+                            rhyme_score = 0.6
             
                 # 2. Syllable match (25%)
                 target_syl = pronouncing.syllable_count(target_phones)
@@ -94,7 +143,8 @@ class ParodyWordSuggestionTool(Tool):
                         "string_similarity": round(string_similarity, 3),
                         "syllables": word_syl,
                         "phones": word_phones,
-                        "rhyme_part": " ".join(word_rhyme_part)
+                        "last_vowel": word_vowel,
+                        "ending": " ".join(word_end) if word_end else ""
                     })
     
         # Sort by similarity score descending
@@ -104,13 +154,10 @@ class ParodyWordSuggestionTool(Tool):
             "target": target,
             "target_syllables": pronouncing.syllable_count(target_phones),
             "target_phones": target_phones,
-            "target_rhyme_part": " ".join(target_rhyme_part),
+            "target_last_vowel": target_vowel,
+            "target_ending": " ".join(target_end) if target_end else "",
             "suggestions": suggestions
         }
     
         return json.dumps(result, indent=2)
 
-
-    def __init__(self, *args, **kwargs):
-        self.is_initialized = False
-