カスタムチャットテンプレートと出力パース入門

from transformers import AutoTokenizer
 
tokenizer = AutoTokenizer.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    use_auth_token="<your-hf-token>"
)
chat_template = tokenizer.chat_template

from jinja2 import Environment, TemplateSyntaxError
 
custom_template = """
{{ bos_token }}
{%- for message in messages %}
    {%- if message['role'] == 'user' %}
        {{ '<|User|>' + message['content'] }}
    {%- elif message['role'] == 'assistant' %}
        {{ '<|Assistant|>' + message['content'] + eos_token }}
    {%- endif %}
{%- endfor %}
{% if add_generation_prompt %}{{ '<|Assistant|>' }}{% endif %}
"""
 
# Validate syntax before use
try:
    Environment().parse(custom_template)
except TemplateSyntaxError as e:
    raise ValueError(f"Invalid Jinja template at line {e.lineno}: {e.message}")

from jinja2 import Template
from datetime import datetime
 
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "What is the population of Bogota?"}
]
 
tools = [
    {
        "type": "function",
        "function": {
            "name": "get_population",
            "description": "Returns the population of a city.",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "Name of the city"}
                },
                "required": ["city"]
            }
        }
    }
]
 
context = {
    "messages": messages,
    "tools": tools,
    "add_generation_prompt": True,
    "bos_token": "<|begin_of_text|>",
    "eos_token": "<|eot_id|>",
    "date_string": datetime.now().strftime("%d %b %Y"),
}
 
rendered_prompt = Template(chat_template).render(**context).strip()

from sambanova import SambaNova
 
client = SambaNova(
    api_key="<your-sambanova-api-key>",
    base_url="https://api.sambanova.ai/v1"
)
 
response = client.completions.create(
    model="Meta-Llama-3.1-8B-Instruct",
    prompt=rendered_prompt,
    max_tokens=2048,
    temperature=0.0,
)
 
raw_output = response.choices[0].text

import json
 
def parse_llama_output(response: str) -> list[dict]:
    """Extract JSON tool calls from a Llama-style response."""
    tool_calls = []
    brace_count, start = 0, None
 
    for i, ch in enumerate(response):
        if ch == "{":
            if brace_count == 0:
                start = i
            brace_count += 1
        elif ch == "}" and start is not None:
            brace_count -= 1
            if brace_count == 0:
                block = response[start:i + 1]
                obj = json.loads(block)
                tool_calls.append({
                    "type": "function",
                    "function": {
                        "name": obj["name"],
                        "arguments": json.dumps(obj["parameters"])
                    }
                })
                start = None
    return tool_calls

import re
import json
 
def parse_deepseek_output(response: str) -> list[dict]:
    """Extract XML-delimited tool calls from a DeepSeek-style response."""
    tool_calls = []
    pattern = r"<｜tool▁call▁begin｜>(.*?)<｜tool▁sep｜>(.*?)<｜tool▁call▁end｜>"
 
    for name, args in re.findall(pattern, response, re.DOTALL):
        tool_calls.append({
            "type": "function",
            "function": {
                "name": name.strip(),
                "arguments": args.strip()
            }
        })
    return tool_calls

def build_assistant_message(response: str, tool_calls: list) -> dict:
    if tool_calls:
        return {"role": "assistant", "content": None, "tool_calls": tool_calls}
    return {"role": "assistant", "content": response.strip(), "tool_calls": []}

def parse(response: str) -> list[dict]:
    """
    Custom parser template.
    Returns a list of tool-call dicts:
    [{"type": "function", "function": {"name": str, "arguments": str}}, ...]
    """
    # Implement your extraction logic here
    return []