Upload tool

tool.py

@@ -1,19 +1,17 @@
 from smolagents.tools import Tool
-import string
 import json
-import difflib
 import pronouncing
+import string
 
 class ParodyWordSuggestionTool(Tool):
     name = "parody_word_suggester"
-    description = """Suggests rhyming funny words using CMU dictionary pronunciations.
+    description = """Suggests rhyming funny words using CMU dictionary and custom 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)', '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}, 'custom_phones': {'type': 'object', 'description': 'Optional dictionary of custom word pronunciations', 'nullable': True}}
     output_type = "string"
     VOWEL_REF = "AH,UH,AX|AE,EH|IY,IH|AO,AA|UW,UH|AY,EY|OW,AO|AW,AO|OY,OW|ER,AXR"
 
     def _get_vowel_groups(self):
-        """Convert the simple string format to usable groups."""
         groups = []
         group_strs = self.VOWEL_REF.split("|")
         for group_str in group_strs:
@@ -21,15 +19,23 @@ 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_last_syllable(self, phones: list) -> tuple:
         """Extract the last syllable (vowel + remaining consonants)."""
         last_vowel_idx = -1
         last_vowel = None
         vowel_groups = self._get_vowel_groups()
     
-        # Find the last vowel
         for i, phone in enumerate(phones):
-            # Strip stress markers for checking
             base_phone = phone.rstrip('012')
             for group in vowel_groups:
                 if base_phone in group:
@@ -40,14 +46,11 @@ class ParodyWordSuggestionTool(Tool):
         if last_vowel_idx == -1:
             return None, []
         
-        # Get all consonants after the vowel
         remaining = phones[last_vowel_idx + 1:]
-    
         return last_vowel, remaining
 
 
     def _strip_stress(self, phones: list) -> list:
-        """Remove stress markers from phones."""
         result = []
         for phone in phones:
             result.append(phone.rstrip('012'))
@@ -55,7 +58,6 @@ class ParodyWordSuggestionTool(Tool):
 
 
     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')
     
@@ -69,14 +71,87 @@ class ParodyWordSuggestionTool(Tool):
         return False
 
 
-    def forward(self, target: str, word_list_str: str, min_similarity: str = "0.5") -> str:
-        """Get rhyming word suggestions."""
+    def _calculate_similarity(self, word1, phones1, word2, phones2):
+        """Calculate similarity score using improved metrics."""
+        # Initialize all variables
+        word_vowel = None
+        word_end = []
+        target_vowel = None
+        target_end = []
+        phone_diff = 0
+        max_phones = 0
+        length_score = 0.0
+        rhyme_score = 0.0
+        stress_score = 0.0
+        i = 0  # For loop counter
+        word_end_clean = []
+        target_end_clean = []
+        matched = 0
+        common_length = 0
+    
+        phone_list1 = phones1.split()
+        phone_list2 = phones2.split()
+    
+        # Calculate length similarity score
+        phone_diff = abs(len(phone_list1) - len(phone_list2))
+        max_phones = max(len(phone_list1), len(phone_list2))
+        length_score = 1.0 if phone_diff == 0 else 1.0 - (phone_diff / max_phones)
+    
+        # Get last syllable components
+        result1 = self._get_last_syllable(phone_list1)
+        result2 = self._get_last_syllable(phone_list2)
+        word_vowel, word_end = result1
+        target_vowel, target_end = result2
+    
+        # Calculate rhyme score
+        rhyme_score = 0.0
+        if word_vowel and target_vowel:
+            if self._vowels_match(word_vowel, target_vowel):
+                word_end_clean = self._strip_stress(word_end)
+                target_end_clean = self._strip_stress(target_end)
+            
+                if word_end_clean == target_end_clean:
+                    rhyme_score = 1.0
+                else:
+                    # Partial rhyme based on ending similarity
+                    common_length = min(len(word_end_clean), len(target_end_clean))
+                    matched = 0
+                    for i in range(common_length):
+                        if word_end_clean[i] == target_end_clean[i]:
+                            matched += 1
+                    rhyme_score = 0.6 * (matched / max(len(word_end_clean), len(target_end_clean)))
+    
+        # Calculate stress pattern similarity
+        import pronouncing
+        stress1 = pronouncing.stresses(phones1)
+        stress2 = pronouncing.stresses(phones2)
+        stress_score = 1.0 if stress1 == stress2 else 0.5
+    
+        # Weighted combination (60% rhyme, 30% length, 10% stress)
+        similarity = (
+            (rhyme_score * 0.6) +
+            (length_score * 0.3) +
+            (stress_score * 0.1)
+        )
+    
+        # Cap at 1.0
+        similarity = min(1.0, similarity)
+    
+        return {
+            "similarity": round(similarity, 3),
+            "rhyme_score": round(rhyme_score, 3),
+            "length_score": round(length_score, 3),
+            "stress_score": round(stress_score, 3),
+            "phone_length_difference": phone_diff
+        }
+
+
+    def forward(self, target: str, word_list_str: str, min_similarity: str = "0.5", custom_phones: dict = None) -> str:
         import pronouncing
         import string
         import json
-        from difflib import SequenceMatcher
     
-        # Initialize variables
+        # Initialize all variables
         target = target.lower().strip(string.punctuation)
         min_similarity = float(min_similarity)
         suggestions = []
@@ -84,6 +159,9 @@ class ParodyWordSuggestionTool(Tool):
         word_end = []
         target_vowel = None
         target_end = []
+        valid_words = []
+        invalid_words = []
+        target_phone_list = []
     
         # Parse JSON string to list
         try:
@@ -95,92 +173,54 @@ class ParodyWordSuggestionTool(Tool):
             }, indent=2)
         
         # Get target pronunciation
-        target_phones = pronouncing.phones_for_word(target)
+        target_phones = self._get_word_phones(target, custom_phones)
         if not target_phones:
             return json.dumps({
-                "error": f"Target word '{target}' not found in CMU dictionary",
+                "error": f"Target word '{target}' not found in dictionary or custom phones",
                 "suggestions": []
             }, indent=2)
     
-        # Filter word list to only words in CMU dictionary
+        # Filter word list
         valid_words = []
         invalid_words = []
         for word in words:
             word = word.lower().strip(string.punctuation)
-            if pronouncing.phones_for_word(word):
+            if self._get_word_phones(word, custom_phones):
                 valid_words.append(word)
             else:
                 invalid_words.append(word)
     
         if not valid_words:
             return json.dumps({
-                "error": "No valid words found in CMU dictionary",
+                "error": "No valid words found in dictionary or custom phones",
                 "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
         for word in valid_words:
-            phones = pronouncing.phones_for_word(word)
-            if phones:
-                word_phones = phones[0]
-                word_phone_list = word_phones.split()
-                word_vowel, word_end = self._get_last_syllable(word_phone_list)
-            
-                # 1. Rhyme score (most important - 60%)
-                rhyme_score = 0.0
-                if word_vowel and target_vowel:
-                    # Check if the vowels are similar
-                    if self._vowels_match(word_vowel, target_vowel):
-                        # Check if endings match (ignoring stress numbers)
-                        word_end_clean = self._strip_stress(word_end)
-                        target_end_clean = self._strip_stress(target_end)
-                    
-                        if word_end_clean == target_end_clean:
-                            rhyme_score = 1.0
-                            # Extra boost for exact match minus first letter
-                            if len(word) == len(target) and word[1:] == target[1:]:
-                                rhyme_score = 1.2
-                        else:
-                            rhyme_score = 0.6
-            
-                # 2. Syllable match (25%)
-                target_syl = pronouncing.syllable_count(target_phones)
-                word_syl = pronouncing.syllable_count(word_phones)
-                syllable_score = 1.0 if target_syl == word_syl else 0.0
-            
-                # 3. String similarity (15%)
-                # Higher weight for end of word similarity
-                if len(word) > 1 and len(target) > 1:
-                    end_similarity = SequenceMatcher(None, word[1:], target[1:]).ratio()
-                    string_similarity = end_similarity
-                else:
-                    string_similarity = SequenceMatcher(None, target, word).ratio()
-            
-                # Combined score
-                similarity = (rhyme_score * 0.6) + (syllable_score * 0.25) + (string_similarity * 0.15)
+            word_phones = self._get_word_phones(word, custom_phones)
+            if word_phones:
+                similarity_result = self._calculate_similarity(word, word_phones, target, target_phones)
             
-                if similarity >= min_similarity:
+                if similarity_result["similarity"] >= min_similarity:
+                    word_phone_list = word_phones.split()
+                    word_vowel, word_end = self._get_last_syllable(word_phone_list)
+                
                     suggestions.append({
                         "word": word,
-                        "similarity": round(similarity, 3),
-                        "rhyme_match": rhyme_score > 0,
-                        "rhyme_score": round(rhyme_score, 3),
-                        "syllable_match": syllable_score == 1.0,
-                        "string_similarity": round(string_similarity, 3),
-                        "syllables": word_syl,
+                        "similarity": similarity_result["similarity"],
+                        "rhyme_score": similarity_result["rhyme_score"],
+                        "length_score": similarity_result["length_score"],
+                        "stress_score": similarity_result["stress_score"],
+                        "phone_length_difference": similarity_result["phone_length_difference"],
                         "phones": word_phones,
                         "last_vowel": word_vowel,
                         "ending": " ".join(word_end) if word_end else "",
-                        "debug_info": {
-                            "word_end_clean": word_end_clean,
-                            "target_end_clean": target_end_clean,
-                            "exact_match": word_end_clean == target_end_clean
-                        }
+                        "is_custom": word in custom_phones if custom_phones else False
                     })
     
         # Sort by similarity score descending
@@ -188,11 +228,10 @@ class ParodyWordSuggestionTool(Tool):
     
         result = {
             "target": target,
-            "target_syllables": pronouncing.syllable_count(target_phones),
             "target_phones": target_phones,
             "target_last_vowel": target_vowel,
             "target_ending": " ".join(target_end) if target_end else "",
-            "invalid_words": invalid_words,  # List of words not in CMU
+            "invalid_words": invalid_words,
             "suggestions": suggestions
         }