Upload tool

requirements.txt

@@ -1,2 +1,2 @@
-pronouncing
 smolagents
+pronouncing

tool.py

@@ -1,13 +1,14 @@
 from smolagents.tools import Tool
-import string
 import pronouncing
 import json
+import string
+import difflib
 
 class ParodyWordSuggestionTool(Tool):
     name = "parody_word_suggester"
-    description = """Suggests rhyming funny words using CMU dictionary and custom pronunciations.
+    description = """Suggests rhyming funny words using CMU dictionary pronunciations.
     Returns similar-sounding words that rhyme, especially focusing on common vowel sounds."""
-    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)', 'default': '0.5', 'nullable': True}, 'custom_phones': {'type': 'object', 'description': 'Optional dictionary of custom word pronunciations', 'default': None, 'nullable': True}}
+    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, 'default': '0.5'}}
     output_type = "string"
     VOWEL_REF = "AH,AX|UH|AE,EH|IY,IH|AO,AA|UW|AY,EY|OW,AO|AW,AO|OY,OW|ER,AXR"
     CONSONANT_REF = "M,N,NG|P,B|T,D|K,G|F,V|TH,DH|S,Z|SH,ZH|L,R|W,Y"
@@ -20,14 +21,12 @@ class ParodyWordSuggestionTool(Tool):
         return groups
 
 
-    def _get_word_phones(self, word, custom_phones=None):
-        """Get phones for a word, checking custom dictionary first."""
-        if custom_phones and word in custom_phones:
-            return custom_phones[word]["primary_phones"]
-        
-        import pronouncing
-        phones = pronouncing.phones_for_word(word)
-        return phones[0] if phones else None
+    def _get_consonant_groups(self):
+        groups = []
+        group_strs = self.CONSONANT_REF.split("|")
+        for group_str in group_strs:
+            groups.append(group_str.split(","))
+        return groups
 
 
     def _get_last_syllable(self, phones: list) -> tuple:
@@ -72,132 +71,109 @@ class ParodyWordSuggestionTool(Tool):
         return False
 
 
-    def _strip_common_suffix(self, phones: list) -> tuple:
-        """Strip common suffixes and return base and suffix phones."""
-        # Initialize variables
-        suffix_name = ""
-        suffix_phones = []
-        phone1 = ""
-        phone2 = ""
-    
-        # Common suffix patterns in CMU phonetic representation
-        SUFFIXES = {
-            'ING': ['IH0', 'NG'],  # -ing
-            'ED': ['EH0', 'D'],    # -ed
-            'ER': ['ER0'],         # -er
-            'EST': ['EH0', 'S', 'T'],  # -est
-            'LY': ['L', 'IY0'],    # -ly
-            'NESS': ['N', 'EH0', 'S'],  # -ness
-        }
+    def _calculate_similarity(self, word1, phones1, word2, phones2):
+        """Calculate similarity with heavy emphasis on rhyming."""
+        from difflib import SequenceMatcher
+        import pronouncing
     
-        for suffix_name, suffix_phones in SUFFIXES.items():
-            if len(phones) > len(suffix_phones):
-                if all(phone1.rstrip('012') == phone2.rstrip('012') 
-                      for phone1, phone2 in zip(phones[-len(suffix_phones):], suffix_phones)):
-                    return phones[:-len(suffix_phones)], suffix_phones
+        # Initialize all variables
+        rhyme_score = 0.0
+        string_score = 0.0
+        pattern_score = 0.0
+        phone_list1 = []
+        phone_list2 = []
+        vowel1 = None
+        vowel2 = None
+        end1 = []
+        end2 = []
+        end1_clean = []
+        end2_clean = []
+        matching_consonants = 0
     
-        return phones, []
-
-
-    def _get_consonant_groups(self):
-        """Get consonant groups from reference string."""
-        groups = []
-        group_strs = self.CONSONANT_REF.split("|")
-        for group_str in group_strs:
-            groups.append(group_str.split(","))
-        return groups
-
-
-    def _consonants_similarity(self, c1: str, c2: str) -> float:
-        """Calculate similarity score between two consonants."""
-        if c1 == c2:
-            return 1.0
-        
-        # Check if they're in the same group
-        consonant_groups = self._get_consonant_groups()
-        for group in consonant_groups:
-            if c1 in group and c2 in group:
-                # Nasals (first group) are more similar to each other
-                if group == consonant_groups[0]:  # M,N,NG group
-                    return 0.8
-                return 0.5
-            
-        return 0.0
-
-
-    def _calculate_similarity(self, word1, phones1, word2, phones2):
-        """Calculate similarity score with enhanced consonant matching."""
-        # Initialize variables as before
         phone_list1 = phones1.split()
         phone_list2 = phones2.split()
     
-        # Get stressed vowels and their positions
-        vowel_idx1 = -1
-        vowel_idx2 = -1
-        primary_vowel1 = None
-        primary_vowel2 = None
+        # Get last syllables
+        vowel1, end1 = self._get_last_syllable(phone_list1)
+        vowel2, end2 = self._get_last_syllable(phone_list2)
     
-        for i, phone in enumerate(phone_list1):
-            if '1' in phone:  # Primary stress
-                vowel_idx1 = i
-                primary_vowel1 = phone.rstrip('012')
-                break
+        # Calculate rhyme score (60%)
+        if vowel1 and vowel2:
+            # Perfect vowel match
+            if vowel1.rstrip('012') == vowel2.rstrip('012'):
+                rhyme_score = 1.0
+            # Similar vowel match
+            elif self._vowels_match(vowel1, vowel2):
+                rhyme_score = 0.8
             
-        for i, phone in enumerate(phone_list2):
-            if '1' in phone:
-                vowel_idx2 = i
-                primary_vowel2 = phone.rstrip('012')
-                break
-
-        # Calculate vowel similarity (50% of total score)
-        vowel_score = 0.0
-        if primary_vowel1 and primary_vowel2:
-            if primary_vowel1 == primary_vowel2:
-                vowel_score = 1.0
-            elif self._vowels_match(primary_vowel1, primary_vowel2):
-                vowel_score = 0.8
-
-        # Calculate consonant similarity (30% of total score)
-        consonant_score = 0.0
-        if vowel_idx1 >= 0 and vowel_idx2 >= 0:
-            # Compare consonants around the stressed vowel
-            pre_c1 = phone_list1[vowel_idx1-1] if vowel_idx1 > 0 else None
-            pre_c2 = phone_list2[vowel_idx2-1] if vowel_idx2 > 0 else None
-            post_c1 = phone_list1[vowel_idx1+1] if vowel_idx1 < len(phone_list1)-1 else None
-            post_c2 = phone_list2[vowel_idx2+1] if vowel_idx2 < len(phone_list2)-1 else None
-        
-            if pre_c1 and pre_c2:
-                consonant_score += self._consonants_similarity(pre_c1, pre_c2)
-            if post_c1 and post_c2:
-                consonant_score += self._consonants_similarity(post_c1, post_c2)
+            # Check endings
+            if end1 and end2:
+                end1_clean = self._strip_stress(end1)
+                end2_clean = self._strip_stress(end2)
             
-            consonant_score = consonant_score / 2  # Normalize to 0-1
-
-        # Pattern/length similarity (20% of total score)
-        pattern_score = 0.0
+                # Perfect ending match
+                if end1_clean == end2_clean:
+                    rhyme_score = min(1.0, rhyme_score + 0.2)
+                # Partial ending match
+                else:
+                    consonant_groups = self._get_consonant_groups()
+                    matching_consonants = 0
+                    for c1, c2 in zip(end1_clean, end2_clean):
+                        if c1 == c2:
+                            matching_consonants += 1
+                        else:
+                            # Check if consonants are in same group
+                            for group in consonant_groups:
+                                if c1 in group and c2 in group:
+                                    matching_consonants += 0.5
+                                    break
+                
+                    if matching_consonants > 0:
+                        rhyme_score = min(1.0, rhyme_score + (0.1 * matching_consonants))
+    
+        # String similarity (25%)
+        if len(word1) > 1 and len(word2) > 1:
+            end_similarity = SequenceMatcher(None, word1[1:], word2[1:]).ratio()
+            string_score = end_similarity
+        else:
+            string_score = SequenceMatcher(None, word1, word2).ratio()
+    
+        # Pattern/Length score (15%)
         if len(phone_list1) == len(phone_list2):
             pattern_score = 1.0
         else:
             pattern_score = 1.0 - (abs(len(phone_list1) - len(phone_list2)) / max(len(phone_list1), len(phone_list2)))
-
-        # Combined weighted score
+    
+        # Final weighted score
         similarity = (
-            (vowel_score * 0.5) +      # Vowel similarity most important
-            (consonant_score * 0.3) +   # Consonant similarity next
-            (pattern_score * 0.2)       # Pattern/length least important
+            (rhyme_score * 0.60) +
+            (string_score * 0.25) +
+            (pattern_score * 0.15)
         )
-
+    
+        # Extra boost for exact matches minus first letter
+        if len(word1) == len(word2) and word1[1:] == word2[1:]:
+            similarity = min(1.0, similarity * 1.2)
+    
+        # Extra penalty for very different lengths
+        if abs(len(word1) - len(word2)) > 2:
+            similarity *= 0.7
+    
         return {
             "similarity": round(similarity, 3),
-            "vowel_score": round(vowel_score, 3),
-            "consonant_score": round(consonant_score, 3),
+            "rhyme_score": round(rhyme_score, 3),
+            "string_score": round(string_score, 3),
             "pattern_score": round(pattern_score, 3),
-            "primary_vowels": f"{primary_vowel1}-{primary_vowel2}",
-            "consonants": "similar" if consonant_score > 0.5 else "different"
+            "details": {
+                "last_vowel_match": vowel1.rstrip('012') == vowel2.rstrip('012') if vowel1 and vowel2 else False,
+                "similar_vowels": self._vowels_match(vowel1, vowel2) if vowel1 and vowel2 else False,
+                "ending_match": " ".join(end1_clean) == " ".join(end2_clean) if end1 and end2 else False,
+                "string_length_diff": abs(len(word1) - len(word2))
+            }
         }
 
 
-    def forward(self, target: str, word_list_str: str, min_similarity: str = "0.5", custom_phones: dict = None) -> str:
+    def forward(self, target: str, word_list_str: str, min_similarity: str = "0.5") -> str:
         import pronouncing
         import string
         import json
@@ -230,38 +206,39 @@ class ParodyWordSuggestionTool(Tool):
             }, indent=2)
         
         # Get target pronunciation
-        target_phones = self._get_word_phones(target, custom_phones)
+        target_phones = pronouncing.phones_for_word(target)
         if not target_phones:
             return json.dumps({
-                "error": f"Target word '{target}' not found in dictionary or custom phones",
+                "error": f"Target word '{target}' not found in CMU dictionary",
                 "suggestions": []
             }, indent=2)
     
-        # Filter word list
+        # Filter word list to only words in CMU dictionary
         valid_words = []
         invalid_words = []
         for word in words:
             word = word.lower().strip(string.punctuation)
-            if self._get_word_phones(word, custom_phones):
+            if pronouncing.phones_for_word(word):
                 valid_words.append(word)
             else:
                 invalid_words.append(word)
     
         if not valid_words:
             return json.dumps({
-                "error": "No valid words found in dictionary or custom phones",
+                "error": "No valid words found in CMU dictionary",
                 "invalid_words": invalid_words,
                 "suggestions": []
             }, indent=2)
     
+        target_phones = target_phones[0]
         target_phone_list = target_phones.split()
         target_vowel, target_end = self._get_last_syllable(target_phone_list)
     
-        # Check each word
         # Check each word
         for word in valid_words:
-            word_phones = self._get_word_phones(word, custom_phones)
-            if word_phones:
+            phones = pronouncing.phones_for_word(word)
+            if phones:
+                word_phones = phones[0]
                 similarity_result = self._calculate_similarity(word, word_phones, target, target_phones)
             
                 if similarity_result["similarity"] >= min_similarity:
@@ -271,15 +248,13 @@ class ParodyWordSuggestionTool(Tool):
                     suggestions.append({
                         "word": word,
                         "similarity": similarity_result["similarity"],
-                        "vowel_score": similarity_result["vowel_score"],
-                        "consonant_score": similarity_result["consonant_score"],
+                        "rhyme_score": similarity_result["rhyme_score"],
+                        "string_score": similarity_result["string_score"],
                         "pattern_score": similarity_result["pattern_score"],
-                        "primary_vowels": similarity_result["primary_vowels"],
-                        "consonants": similarity_result["consonants"],
                         "phones": word_phones,
                         "last_vowel": word_vowel,
                         "ending": " ".join(word_end) if word_end else "",
-                        "is_custom": word in custom_phones if custom_phones else False
+                        "details": similarity_result["details"]
                     })
     
         # Sort by similarity score descending