adding pydantic sampler for use with strict mode

Author: gittb

backends/exllamav2/grammar.py

@@ -7,8 +7,9 @@
     build_token_enforcer_tokenizer_data,
 )
 from loguru import logger
-from typing import List
+from typing import List, Type
 from functools import lru_cache
+from pydantic import BaseModel
 
 
 class OutlinesTokenizerWrapper:
@@ -106,6 +107,34 @@ def add_json_schema_filter(
         # Append the filters
         self.filters.extend([lmfilter, prefix_filter])
 
+    def add_pydantic_filter(
+            self,
+            pydantic_model: Type[BaseModel],
+            model: ExLlamaV2,
+            tokenizer:ExLlamaV2Tokenizer
+    ):
+        """Adds an ExllamaV2 filter based on a Pydantic model"""
+        # Create the parser
+        try:
+            schema_parser = JsonSchemaParser(pydantic_model.model_json_schema())
+        except Exception:
+            traceback.print_exc()
+            logger.error(
+                "Skipping because the pydantic model couldn't be used. "
+                "Please read the above error for more information."
+            )
+
+            return
+        
+        json_prefixes = ["[", "{"]
+
+        lmfilter = ExLlamaV2TokenEnforcerFilter(
+            schema_parser, _get_lmfe_tokenizer_data(tokenizer)
+        )
+        prefix_filter = ExLlamaV2PrefixFilter(model, tokenizer, json_prefixes)
+
+        self.filters.extend([lmfilter, prefix_filter])
+
     def add_regex_filter(
         self,
         pattern: str,

backends/exllamav2/model.py

@@ -1077,6 +1077,13 @@ async def generate_gen(
                 json_schema, self.model, self.tokenizer
             )
 
+        # Add pydantic filter if it exists
+        pydantic_model = unwrap(kwargs.get("pydantic_model"))
+        if pydantic_model:
+            grammar_handler.add_pydantic_filter(
+                pydantic_model, self.model, self.tokenizer
+            )
+
         # Add regex filter if it exists
         regex_pattern = unwrap(kwargs.get("regex_pattern"))
         if regex_pattern:

common/sampling.py

@@ -5,7 +5,7 @@
 from copy import deepcopy
 from loguru import logger
 from pydantic import AliasChoices, BaseModel, Field
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union, Type
 
 from common.utils import unwrap, prune_dict
 
@@ -176,6 +176,10 @@ class BaseSamplerRequest(BaseModel):
         default_factory=lambda: get_default_sampler_value("json_schema"),
     )
 
+    pydantic_model: Optional[Type[BaseModel]] = Field(
+        default_factory=lambda: get_default_sampler_value("pydantic_model"),
+    )
+
     regex_pattern: Optional[str] = Field(
         default_factory=lambda: get_default_sampler_value("regex_pattern"),
     )
@@ -329,6 +333,7 @@ def to_gen_params(self, **kwargs):
             "cfg_scale": self.cfg_scale,
             "negative_prompt": self.negative_prompt,
             "json_schema": self.json_schema,
+            "pydantic_model": self.pydantic_model,
             "regex_pattern": self.regex_pattern,
             "grammar_string": self.grammar_string,
             "speculative_ngram": self.speculative_ngram,