More ReDOS fixes! (#36964)

* More ReDOS fixes!

* Slight regex cleanup

* Cleanup regex replacement

* Drop that regex entirely too

* The regex didn't match config.json, let's make sure we don't either

* Cleanup allowed_value_chars a little

* Cleanup the import search

* Catch multi-condition blocks too

* Trigger tests

* Trigger tests

Author: Rocketknight1

src/transformers/commands/chat.py

@@ -17,7 +17,7 @@
 import json
 import os
 import platform
-import re
+import string
 import time
 from argparse import ArgumentParser, Namespace
 from dataclasses import dataclass, field
@@ -44,6 +44,10 @@
 
     from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
 
+ALLOWED_KEY_CHARS = set(string.ascii_letters + string.whitespace)
+ALLOWED_VALUE_CHARS = set(
+    string.ascii_letters + string.digits + string.whitespace + r".!\"#$%&'()*+,\-/:<=>?@[]^_`{|}~"
+)
 
 HELP_STRING = """\
 
@@ -71,8 +75,6 @@
     "repetition_penalty",
 ]
 
-SETTING_RE = r"^set\s+[A-Za-z\s_]+=[A-Za-z\d\s.!\"#$%&'()*+,-/:<=>?@\[\]^_`{|}~]+(?:;\s*[A-Za-z\s_]+=[A-Za-z\d\s.!\"#$%&'()*+,-/:<=>?@\[\]^_`{|}~]+)*$"
-
 DEFAULT_EXAMPLES = {
     "llama": {"text": "There is a Llama in my lawn, how can I get rid of it?"},
     "code": {
@@ -438,6 +440,36 @@ def register_subcommand(parser: ArgumentParser):
     def __init__(self, args):
         self.args = args
 
+    @staticmethod
+    def is_valid_setting_command(s: str) -> bool:
+        # First check the basic structure
+        if not s.startswith("set ") or "=" not in s:
+            return False
+
+        # Split into individual assignments
+        assignments = [a.strip() for a in s[4:].split(";") if a.strip()]
+
+        for assignment in assignments:
+            # Each assignment should have exactly one '='
+            if assignment.count("=") != 1:
+                return False
+
+            key, value = assignment.split("=", 1)
+            key = key.strip()
+            value = value.strip()
+            if not key or not value:
+                return False
+
+            # Keys can only have alphabetic characters, spaces and underscores
+            if not set(key).issubset(ALLOWED_KEY_CHARS):
+                return False
+
+            # Values can have just about anything that isn't a semicolon
+            if not set(value).issubset(ALLOWED_VALUE_CHARS):
+                return False
+
+        return True
+
     def run(self):
         if not is_rich_available():
             raise ImportError("You need to install rich to use the chat interface. (`pip install rich`)")
@@ -499,7 +531,7 @@ def run(self):
                     interface.print_green(f"Chat saved in {filename}!")
                     continue
 
-                if re.match(SETTING_RE, user_input):
+                if self.is_valid_setting_command(user_input):
                     current_args, success = parse_settings(user_input, current_args, interface)
                     if success:
                         chat = []

src/transformers/configuration_utils.py

@@ -17,7 +17,6 @@
 import copy
 import json
 import os
-import re
 import warnings
 from typing import Any, Optional, Union
 
@@ -44,8 +43,6 @@
 
 logger = logging.get_logger(__name__)
 
-_re_configuration_file = re.compile(r"config\.(.*)\.json")
-
 
 class PretrainedConfig(PushToHubMixin):
     # no-format
@@ -1160,9 +1157,8 @@ def get_configuration_file(configuration_files: list[str]) -> str:
     """
     configuration_files_map = {}
     for file_name in configuration_files:
-        search = _re_configuration_file.search(file_name)
-        if search is not None:
-            v = search.groups()[0]
+        if file_name.startswith("config.") and file_name.endswith(".json") and file_name != "config.json":
+            v = file_name.removeprefix("config.").removesuffix(".json")
             configuration_files_map[v] = file_name
     available_versions = sorted(configuration_files_map.keys())
 

src/transformers/dynamic_module_utils.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Utilities to dynamically load objects from the Hub."""
 
+import ast
 import filecmp
 import hashlib
 import importlib
@@ -148,22 +149,38 @@ def get_imports(filename: Union[str, os.PathLike]) -> list[str]:
     """
     with open(filename, encoding="utf-8") as f:
         content = f.read()
-
-    # filter out try/except block so in custom code we can have try/except imports
-    content = re.sub(r"\s*try\s*:.*?except.*?:", "", content, flags=re.DOTALL)
-
-    # filter out imports under is_flash_attn_2_available block for avoid import issues in cpu only environment
-    content = re.sub(
-        r"if is_flash_attn[a-zA-Z0-9_]+available\(\):\s*(from flash_attn\s*.*\s*)+", "", content, flags=re.MULTILINE
-    )
-
-    # Imports of the form `import xxx`
-    imports = re.findall(r"^\s*import\s+(\S+)\s*$", content, flags=re.MULTILINE)
-    # Imports of the form `from xxx import yyy`
-    imports += re.findall(r"^\s*from\s+(\S+)\s+import", content, flags=re.MULTILINE)
-    # Only keep the top-level module
-    imports = [imp.split(".")[0] for imp in imports if not imp.startswith(".")]
-    return list(set(imports))
+    imported_modules = set()
+
+    def recursive_look_for_imports(node):
+        if isinstance(node, ast.Try):
+            return  #  Don't recurse into Try blocks and ignore imports in them
+        elif isinstance(node, ast.If):
+            test = node.test
+            for condition_node in ast.walk(test):
+                if isinstance(condition_node, ast.Call) and condition_node.func.id.startswith("is_flash_attn"):
+                    # Don't recurse into "if flash_attn_available()" blocks and ignore imports in them
+                    return
+        elif isinstance(node, ast.Import):
+            # Handle 'import x' statements
+            for alias in node.names:
+                top_module = alias.name.split(".")[0]
+                if top_module:
+                    imported_modules.add(top_module)
+        elif isinstance(node, ast.ImportFrom):
+            # Handle 'from x import y' statements, ignoring relative imports
+            if node.level == 0 and node.module:
+                top_module = node.module.split(".")[0]
+                if top_module:
+                    imported_modules.add(top_module)
+
+        # Recursively visit all children
+        for child in ast.iter_child_nodes(node):
+            recursive_look_for_imports(child)
+
+    tree = ast.parse(content)
+    recursive_look_for_imports(tree)
+
+    return sorted(imported_modules)
 
 
 def check_imports(filename: Union[str, os.PathLike]) -> list[str]: