dnrops
/
transformers
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Clean up enum regex

This commit is contained in:
Matt 2024-05-28 16:21:22 +01:00
parent c2404ed744
commit 82d90244dd
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/transformers/utils/chat_template_utils.py
View File

@ -301,7 +301,7 @@ def get_json_schema(func):
f"Cannot generate JSON schema for {func.__name__} because the docstring has no description for the argument '{arg}'"
)
desc = param_descriptions[arg]
enum_choices = re.search(r"\(choices:\s*([^)]+)\)\s*$", desc, flags=re.IGNORECASE)
enum_choices = re.search(r"\(choices:\s*(.*?)\)\s*$", desc, flags=re.IGNORECASE)
if enum_choices:
json_schema["properties"][arg]["enum"] = [c.strip() for c in json.loads(enum_choices.group(1))]
desc = enum_choices.string[: enum_choices.start()].strip()