[Bugfix] Migrate to REGEX Library to prevent catastrophic backtracking

Signed-off-by: Crucifixion-Fxl <[email protected]>

Author: Crucifixion-Fxl

.github/scripts/cleanup_pr_body.sh

@@ -26,13 +26,13 @@ sed -i '/\*\*BEFORE SUBMITTING, PLEASE READ.*\*\*/,$d' "${NEW}"
 
 # Remove HTML <details> section that includes <summary> text of "PR Checklist (Click to Expand)"
 python3 - <<EOF
-import re
+import regex
 
 with open("${NEW}", "r") as file:
     content = file.read()
 
-pattern = re.compile(r'(---\n\n)?<details>.*?<summary>.*?PR Checklist \(Click to Expand\).*?</summary>.*?</details>', re.DOTALL)
-content = re.sub(pattern, '', content)
+pattern = regex.compile(r'(---\n\n)?<details>.*?<summary>.*?PR Checklist \(Click to Expand\).*?</summary>.*?</details>', regex.DOTALL)
+content = regex.sub(pattern, '', content)
 
 with open("${NEW}", "w") as file:
     file.write(content)

benchmarks/benchmark_serving_structured_output.py

@@ -672,11 +672,11 @@ def process_one_metric(
 def evaluate(ret, args):
     def _eval_correctness_json(expected, actual):
         # extract json string from string using regex
-        import re
+        import regex
 
         actual = actual.replace("\n", "").replace(" ", "").strip()
         try:
-            actual = re.search(r"\{.*\}", actual).group()
+            actual = regex.search(r"\{.*\}", actual).group()
             actual = json.loads(actual)
         except Exception:
             return False
@@ -687,9 +687,9 @@ def _eval_correctness_choice(expected, actual):
         return actual in args.choice
 
     def _eval_correctness_regex(expected, actual):
-        import re
+        import regex
 
-        return re.match(args.regex, actual) is not None
+        return regex.match(args.regex, actual) is not None
 
     def _eval_correctness(expected, actual):
         if args.structure_type == "guided_json":

benchmarks/kernels/graph_machete_bench.py

@@ -2,11 +2,11 @@
 
 import math
 import pickle
-import re
 from collections import defaultdict
 
 import matplotlib.pyplot as plt
 import pandas as pd
+import regex
 import seaborn as sns
 from torch.utils.benchmark import Measurement as TMeasurement
 
@@ -27,12 +27,12 @@
 
     results = defaultdict(lambda: list())
     for v in raw_results:
-        result = re.search(r"MKN=\(\d+x(\d+x\d+)\)", v.task_spec.sub_label)
+        result = regex.search(r"MKN=\(\d+x(\d+x\d+)\)", v.task_spec.sub_label)
         if result is not None:
             KN = result.group(1)
         else:
             raise Exception("MKN not found")
-        result = re.search(r"MKN=\((\d+)x\d+x\d+\)", v.task_spec.sub_label)
+        result = regex.search(r"MKN=\((\d+)x\d+x\d+\)", v.task_spec.sub_label)
         if result is not None:
             M = result.group(1)
         else:

docs/source/conf.py

@@ -15,10 +15,10 @@
 import datetime
 import logging
 import os
-import re
 import sys
 from pathlib import Path
 
+import regex
 import requests
 
 logger = logging.getLogger(__name__)
@@ -198,7 +198,7 @@ def linkcode_resolve(domain, info):
     for lineno, line in enumerate(lines, 1):
         if not line or line.startswith("#"):
             continue
-        if re.match(pattern, line):
+        if regex.match(pattern, line):
             break
 
     # If the line number is not found, return None

docs/source/generate_examples.py

@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import itertools
-import re
 from dataclasses import dataclass, field
 from pathlib import Path
 
+import regex
+
 ROOT_DIR = Path(__file__).parent.parent.parent.resolve()
 ROOT_DIR_RELATIVE = '../../../..'
 EXAMPLE_DIR = ROOT_DIR / "examples"
@@ -32,7 +33,7 @@ def fix_case(text: str) -> str:
         r"int\d+": lambda x: x.group(0).upper(),  # e.g. int8, int16
     }
     for pattern, repl in subs.items():
-        text = re.sub(rf'\b{pattern}\b', repl, text, flags=re.IGNORECASE)
+        text = regex.sub(rf'\b{pattern}\b', repl, text, flags=regex.IGNORECASE)
     return text
 
 

examples/offline_inference/prithvi_geospatial_mae.py

@@ -20,12 +20,12 @@
 import argparse
 import datetime
 import os
-import re
 from typing import Union
 
 import albumentations
 import numpy as np
 import rasterio
+import regex
 import torch
 from einops import rearrange
 from terratorch.datamodules import Sen1Floods11NonGeoDataModule
@@ -300,7 +300,7 @@ def load_example(
             location_coords.append(coords)
 
         try:
-            match = re.search(r'(\d{7,8}T\d{6})', file)
+            match = regex.search(r'(\d{7,8}T\d{6})', file)
             if match:
                 year = int(match.group(1)[:4])
                 julian_day = match.group(1).split('T')[0][4:]

requirements/common.txt

@@ -1,3 +1,4 @@
+regex # Replace re for higher-performance regex matching
 cachetools
 psutil
 sentencepiece  # Required for LLaMA tokenizer.

requirements/docs.txt

@@ -8,6 +8,7 @@ myst-parser==3.0.1  # `myst-parser==4.0.1` breaks inline code in titles
 msgspec
 snowballstemmer<3  # https://github.com/snowballstem/snowball/issues/229
 commonmark # Required by sphinx-argparse when using :markdownhelp:
+regex # Replace re for higher-performance regex matching
 
 # Custom autodoc2 is necessary for faster docstring processing
 # see: https://github.com/sphinx-extensions2/sphinx-autodoc2/issues/33#issuecomment-2856386035

setup.py

@@ -5,12 +5,12 @@
 import json
 import logging
 import os
-import re
 import subprocess
 import sys
 from pathlib import Path
 from shutil import which
 
+import regex
 import torch
 from packaging.version import Version, parse
 from setuptools import Extension, setup
@@ -389,8 +389,7 @@ def run(self) -> None:
             # vllm_flash_attn python code:
             # Regex from
             #  `glob.translate('vllm/vllm_flash_attn/**/*.py', recursive=True)`
-            import re
-            compiled_regex = re.compile(
+            compiled_regex = regex.compile(
                 r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py")
             file_members += list(
                 filter(lambda x: compiled_regex.match(x.filename),
@@ -510,7 +509,7 @@ def get_neuronxcc_version():
         content = fp.read()
 
     # Extract the version using a regular expression
-    match = re.search(r"__version__ = '(\S+)'", content)
+    match = regex.search(r"__version__ = '(\S+)'", content)
     if match:
         # Return the version string
         return match.group(1)

tests/entrypoints/llm/test_guided_generate.py

@@ -1,12 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import json
-import re
 import weakref
 from enum import Enum
 
 import jsonschema
 import pytest
+import regex
 from pydantic import BaseModel
 
 from vllm.distributed import cleanup_dist_env_and_memory
@@ -62,7 +62,7 @@ def test_guided_regex(sample_regex, llm, guided_decoding_backend: str,
         generated_text = output.outputs[0].text
         print(generated_text)
         assert generated_text is not None
-        assert re.fullmatch(sample_regex, generated_text) is not None
+        assert regex.fullmatch(sample_regex, generated_text) is not None
         print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
 
 
@@ -479,7 +479,7 @@ def test_guided_number_range_json_completion(llm, guided_decoding_backend: str,
         jsonschema.validate(instance=output_json, schema=sample_output_schema)
         assert 18 <= output_json["age"] <= 99
         assert 0.0 <= output_json["score"] <= 100.0
-        assert (re.fullmatch(r"^\d{5}(-\d{4})?$", output_json["zipcode"])
+        assert (regex.fullmatch(r"^\d{5}(-\d{4})?$", output_json["zipcode"])
                 is not None)
 
 

tests/entrypoints/openai/test_chat.py

@@ -2,13 +2,13 @@
 
 # imports for guided decoding tests
 import json
-import re
 from typing import Optional
 
 import jsonschema
 import openai  # use the official client for correctness check
 import pytest
 import pytest_asyncio
+import regex
 import requests
 import torch
 from openai import BadRequestError, OpenAI
@@ -585,7 +585,7 @@ async def test_guided_regex_chat(client: openai.AsyncOpenAI, sample_regex):
         extra_body=dict(guided_regex=sample_regex))
     ip1 = chat_completion.choices[0].message.content
     assert ip1 is not None
-    assert re.fullmatch(sample_regex, ip1) is not None
+    assert regex.fullmatch(sample_regex, ip1) is not None
 
     messages.append({"role": "assistant", "content": ip1})
     messages.append({"role": "user", "content": "Give me a different one"})
@@ -596,7 +596,7 @@ async def test_guided_regex_chat(client: openai.AsyncOpenAI, sample_regex):
         extra_body=dict(guided_regex=sample_regex))
     ip2 = chat_completion.choices[0].message.content
     assert ip2 is not None
-    assert re.fullmatch(sample_regex, ip2) is not None
+    assert regex.fullmatch(sample_regex, ip2) is not None
     assert ip1 != ip2
 
 

tests/entrypoints/openai/test_completion.py

@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
-
 # imports for guided decoding tests
 import json
-import re
 import shutil
 from tempfile import TemporaryDirectory
 from typing import Optional
@@ -11,6 +9,7 @@
 import openai  # use the official client for correctness check
 import pytest
 import pytest_asyncio
+import regex
 # downloading lora to test lora requests
 from huggingface_hub import snapshot_download
 from openai import BadRequestError
@@ -677,8 +676,8 @@ async def test_guided_regex_completion(client: openai.AsyncOpenAI,
     assert completion.id is not None
     assert len(completion.choices) == 3
     for i in range(3):
-        assert re.fullmatch(sample_regex,
-                            completion.choices[i].text) is not None
+        assert regex.fullmatch(sample_regex,
+                               completion.choices[i].text) is not None
 
 
 @pytest.mark.asyncio
@@ -747,7 +746,7 @@ async def test_echo_logprob_completion(client: openai.AsyncOpenAI,
 
         prompt_text = tokenizer.decode(prompt) if isinstance(prompt,
                                                              list) else prompt
-        assert re.search(r"^" + prompt_text, completion.choices[0].text)
+        assert regex.search(r"^" + prompt_text, completion.choices[0].text)
         logprobs = completion.choices[0].logprobs
         assert logprobs is not None
         assert len(logprobs.text_offset) > 5

tests/entrypoints/openai/test_prompt_validation.py

@@ -1,10 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # imports for guided decoding tests
-import re
-
 import openai
 import pytest
+import regex
 
 from ...utils import RemoteOpenAIServer
 
@@ -32,7 +31,7 @@ async def test_out_of_vocab_token_ids():
         client = remote_server.get_async_client()
 
         with pytest.raises(openai.BadRequestError,
-                           match=re.compile('.*out of vocabulary.*')):
+                           match=regex.compile('.*out of vocabulary.*')):
             await client.completions.create(model=model_name,
                                             prompt=[999999],
                                             max_tokens=5,
@@ -47,7 +46,7 @@ async def test_reject_multistep_with_guided_decoding():
         client = remote_server.get_async_client()
 
         with pytest.raises(openai.BadRequestError,
-                           match=re.compile(
+                           match=regex.compile(
                                '.*Guided decoding .* multi-step decoding.*')):
             await client.completions.create(
                 model=model_name,

tests/models/multimodal/generation/test_phi4mm.py

@@ -1,12 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import os
-import re
 from collections.abc import Sequence
 from typing import Optional
 
 import librosa
 import pytest
+import regex
 from huggingface_hub import snapshot_download
 from transformers import AutoTokenizer
 
@@ -44,7 +44,7 @@ def vllm_to_hf_output(vllm_output: tuple[list[int], str,
     """Sanitize vllm output to be comparable with hf output."""
     _, output_str, out_logprobs = vllm_output
 
-    output_str_without_image = re.sub(r"(<\|image_\d+\|>)+", "", output_str)
+    output_str_without_image = regex.sub(r"(<\|image_\d+\|>)+", "", output_str)
     assert output_str_without_image[0] == " "
     output_str_without_image = output_str_without_image[1:]
 

tests/models/multimodal/generation/vlm_utils/model_utils.py

@@ -3,11 +3,11 @@
 for manipulating the input / output of HF & vLLM test runners, which are
 typically specific to a small subset of models.
 """
-import re
 import types
 from pathlib import PosixPath
 from typing import Optional, Union
 
+import regex
 import torch
 from PIL.Image import Image
 from transformers import (AutoConfig, AutoTokenizer, BatchFeature,
@@ -160,7 +160,7 @@ def phi3v_vllm_to_hf_output(vllm_output: RunnerOutput,
     """Sanitize vllm output [phi3v] to be comparable with hf output."""
     _, output_str, out_logprobs = vllm_output
 
-    output_str_without_image = re.sub(r"(<\|image_\d+\|>)+", "", output_str)
+    output_str_without_image = regex.sub(r"(<\|image_\d+\|>)+", "", output_str)
     assert output_str_without_image[0] == " "
     output_str_without_image = output_str_without_image[1:]
 
@@ -335,7 +335,7 @@ def processor(*args, text="", images=None, **kwargs):
 
         images = [images] if isinstance(images, Image) else images
 
-        contents = re.findall(
+        contents = regex.findall(
             r"<\|begin_of_image\|><\|endoftext\|><\|end_of_image\|>(.*?)<\|assistant\|>",
             text,
         )

tests/tool_use/test_tool_choice_required.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 import json
-import re
 from copy import deepcopy
 from unittest.mock import MagicMock
 
@@ -73,7 +72,7 @@ def _compile_and_check(tools: list[ChatCompletionToolsParam], sample_output,
     # use build_regex_from_schema used in JSONLogitsProcessor to create Guide
     from outlines_core.fsm.json_schema import build_regex_from_schema
     regex = build_regex_from_schema(json.dumps(schema))
-    compiled = re.compile(regex)
+    compiled = regex.compile(regex)
     matches = compiled.fullmatch(json.dumps(sample_output)) is not None
 
     assert matches == should_match