Stop TF weight rename reDOS (#38325)

* let's try a non-regex solution

* make fixup

* Slight adjustment

* Let's just use the original code with a check

* slight tweak to conditional

* slight tweak to conditional
This commit is contained in:
Matt 2025-05-26 16:58:51 +01:00 committed by GitHub
parent 0a4e8e2855
commit 701caef704
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -78,6 +78,9 @@ def convert_tf_weight_name_to_pt_weight_name(
tf_name = tf_name[len(name_scope) :]
tf_name = tf_name.lstrip("/")
tf_name = tf_name.replace(":0", "") # device ids
if (len(tf_name) > 2048 and "___" in tf_name) or tf_name.count("___") > 10:
# ReDOS check
raise ValueError("TF variable name is too long or contains too many ___ separators: " + tf_name)
tf_name = re.sub(
r"/[^/]*___([^/]*)/", r"/\1/", tf_name
) # '$1___$2' is replaced by $2 (can be used to duplicate or remove layers in TF2.0 vs PyTorch)