e4cbb28ea1
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Shangming Cai <csmthu@gmail.com>
51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
"""Regenerate the array-only template regression using SGLang's content processor.
|
|
|
|
Run: PYTHONPATH=../../python python tests/scripts/generate_array_content_fixture.py
|
|
"""
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from transformers import PreTrainedTokenizerFast
|
|
|
|
from sglang.srt.parser.jinja_template_utils import (
|
|
detect_jinja_template_content_format,
|
|
process_content_for_template_format,
|
|
)
|
|
|
|
ROOT = Path(__file__).resolve().parents[1] / "fixtures"
|
|
TEMPLATE = (
|
|
"{% for message in messages %}{{ message.role }}:"
|
|
"{% for part in message.content %}"
|
|
"{% if part.type == 'text' %}{{ part.text }}"
|
|
"{% elif part.type == 'image' %}<image>{% endif %}"
|
|
"{% endfor %};{% endfor %}"
|
|
"{% if add_generation_prompt %}assistant:{% endif %}"
|
|
)
|
|
|
|
|
|
def main():
|
|
tokenizer = PreTrainedTokenizerFast(
|
|
tokenizer_file=str(ROOT / "tiny_tokenizer.json")
|
|
)
|
|
tokenizer.chat_template = TEMPLATE
|
|
content_format = detect_jinja_template_content_format(TEMPLATE)
|
|
assert content_format == "openai"
|
|
cases = []
|
|
for content in ["hello", [{"type": "text", "text": "hello"}]]:
|
|
message = process_content_for_template_format(
|
|
{"role": "user", "content": content}, content_format, [], [], [], []
|
|
)
|
|
ids = tokenizer.apply_chat_template(
|
|
[message], return_dict=False, add_generation_prompt=True
|
|
)
|
|
cases.append(json.dumps({"content": content, "engine_token_ids": ids}))
|
|
cases_json = ",\n".join(cases)
|
|
(ROOT / "array_content_rendering.json").write_text(
|
|
f'{{"chat_template": {json.dumps(TEMPLATE)}, "cases": [\n{cases_json}\n]}}\n'
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|