@modular-prompt/driver 0.14.0 → 0.15.0

package/src/pytorch/python/utils/token_utils.py

@@ -0,0 +1,376 @@
+"""
+トークン関連のユーティリティ関数
+"""
+from utils.chat_template_constraints import detect_chat_restrictions
+
+
+def is_eod_token(response, tokenizer):
+    """
+    レスポンスがEODトークンかどうかを判定する
+    
+    Args:
+        response: stream_generateからのレスポンス
+        tokenizer: tokenizerオブジェクト（必須）
+        
+    Returns:
+        bool: EODトークンの場合True
+    """
+    # 1. finish_reasonによる終了判定（MLX-LMの標準的な方法）
+    if hasattr(response, 'finish_reason') and response.finish_reason == 'stop':
+        return True
+    
+    # 2. response.tokenによる終了トークン判定
+    if hasattr(response, 'token'):
+        token = response.token
+        
+        # special_tokens_mapとadded_tokens_encoderから終了トークンを取得
+        end_token_ids = []
+        
+        # special_tokens_mapから標準的な終了トークンを取得
+        if hasattr(tokenizer, 'special_tokens_map') and hasattr(tokenizer, 'added_tokens_encoder'):
+            special_map = tokenizer.special_tokens_map
+            added_encoder = tokenizer.added_tokens_encoder
+            
+            # EOSトークン
+            eos_token_str = special_map.get('eos_token')
+            if eos_token_str and eos_token_str in added_encoder:
+                end_token_ids.append(added_encoder[eos_token_str])
+            
+            # その他の終了関連トークン
+            end_related_keys = ['eoi_token']  # end_of_image
+            for key in end_related_keys:
+                token_str = special_map.get(key)
+                if token_str and token_str in added_encoder:
+                    end_token_ids.append(added_encoder[token_str])
+        
+        # added_tokens_encoderから直接取得（会話終了トークンなど）
+        if hasattr(tokenizer, 'added_tokens_encoder'):
+            added_encoder = tokenizer.added_tokens_encoder
+            conversation_end_tokens = ['<end_of_turn>']
+            for token_str in conversation_end_tokens:
+                token_id = added_encoder.get(token_str)
+                if token_id is not None:
+                    end_token_ids.append(token_id)
+        
+        # 重複を除去してチェック
+        if token in set(end_token_ids):
+            return True
+
+    return False
+
+
+def get_special_tokens(tokenizer):
+    """
+    tokenizerから特殊トークンを取得する
+    
+    Returns:
+        dict: special_tokens情報
+    """
+    special_tokens = {}
+    
+    # 標準的なspecial tokens（tokenizerに定義されているもの）
+    # VLM processorではこれらの属性がない場合があるためgetattr使用
+    standard_tokens = {
+        "eod": getattr(tokenizer, "eos_token", None),  # End of Document/Sequence
+        "bos": getattr(tokenizer, "bos_token", None),  # Beginning of Sequence
+        "unk": getattr(tokenizer, "unk_token", None),  # Unknown token
+        "pad": getattr(tokenizer, "pad_token", None),  # Padding token
+    }
+    
+    for name, token in standard_tokens.items():
+        if token is not None:
+            token_id = getattr(tokenizer, f"{name}_token_id", None)
+            if token_id is not None:
+                special_tokens[name] = {"text": token, "id": token_id}
+    
+    # ペアトークン（存在する場合のみ）
+    pair_tokens = {
+        # ChatML基本形式
+        "system": ("<|system|>", "<|/system|>"),
+        "user": ("<|user|>", "<|/user|>"),
+        "assistant": ("<|assistant|>", "<|/assistant|>"),
+        
+        # フォーマット・構造化
+        "code": ("<|code_start|>", "<|code_end|>"),
+        "python": ("<|python|>", "<|/python|>"),
+        "javascript": ("<|javascript|>", "<|/javascript|>"),
+        "bash": ("<|bash|>", "<|/bash|>"),
+        "quote": ("<|quote|>", "<|/quote|>"),
+        "ref": ("<|ref|>", "<|/ref|>"),
+        "citation": ("<|citation|>", "<|/citation|>"),
+        "table": ("<|table|>", "<|/table|>"),
+        "heading": ("<|heading|>", "<|/heading|>"),
+        
+        # メディア・リッチコンテンツ
+        "image": ("<|image|>", "<|/image|>"),
+        "audio": ("<|audio|>", "<|/audio|>"),
+        "video": ("<|video|>", "<|/video|>"),
+        
+        # 機能・制御
+        "tool_call": ("<|tool_call|>", "<|/tool_call|>"),
+        "function": ("<|function|>", "<|/function|>"),
+        "api": ("<|api|>", "<|/api|>"),
+        "search": ("<|search|>", "<|/search|>"),
+        "knowledge": ("<|knowledge|>", "<|/knowledge|>"),
+        "context": ("<|context|>", "<|/context|>"),
+        
+        # 思考・推論
+        "thinking": ("<|thinking|>", "</thinking>"),
+        "reasoning": ("<|reasoning|>", "<|/reasoning|>"),
+        "scratchpad": ("<|scratchpad|>", "<|/scratchpad|>"),
+        "analysis": ("<|analysis|>", "<|/analysis|>"),
+        "summary": ("<|summary|>", "<|/summary|>"),
+        "explanation": ("<|explanation|>", "<|/explanation|>"),
+
+        # tool_call バリエーション（追加）
+        "tool_call_explicit": ("<|tool_call_start|>", "<|tool_call_end|>"),
+        "tool_call_xml": ("<tool_call>", "</tool_call>"),
+        "tool_calls_section": ("<|tool_calls_section_begin|>", "<|tool_calls_section_end|>"),
+        "function_call_tags": ("<start_function_call>", "<end_function_call>"),
+        "longcat_tool_call": ("<longcat_tool_call>", "</longcat_tool_call>"),
+        "minimax_tool_call": ("<minimax:tool_call>", "</minimax:tool_call>"),
+    }
+    
+    # 単体トークン（存在する場合のみ）
+    single_tokens = {
+        # Fill-in-the-Middle
+        "fim_prefix": "<|fim_prefix|>",
+        "fim_middle": "<|fim_middle|>", 
+        "fim_suffix": "<|fim_suffix|>",
+        
+        # リスト・構造
+        "list_item": "<|list_item|>",
+        
+        # メディア単体
+        "vision": "<|vision|>",
+        
+        # 一般的なマークダウン風
+        "code_inline": "`",
+        "code_block_start": "```",
+        "code_block_end": "```",
+
+        # ツール関連の単体トークン（追加）
+        "tool_calls_marker": "[TOOL_CALLS]",
+        # Harmony形式のcallトークン（tool_call_endとは異なる用途）
+        "harmony_call": "<|call|>",
+    }
+    
+    # VLM processorではconvert_tokens_to_idsがない場合がある
+    if not hasattr(tokenizer, 'convert_tokens_to_ids'):
+        return special_tokens
+
+    unk_token_id = getattr(tokenizer, "unk_token_id", None)
+
+    # ペアトークンの処理
+    for name, (start_token, end_token) in pair_tokens.items():
+        start_id = tokenizer.convert_tokens_to_ids(start_token)
+        end_id = tokenizer.convert_tokens_to_ids(end_token)
+
+        # unk_tokenでない場合のみ追加
+        if start_id != unk_token_id and end_id != unk_token_id:
+            special_tokens[name] = {
+                "start": {"text": start_token, "id": start_id},
+                "end": {"text": end_token, "id": end_id}
+            }
+
+    # 単体トークンの処理
+    for name, token_text in single_tokens.items():
+        token_id = tokenizer.convert_tokens_to_ids(token_text)
+
+        # unk_tokenでない場合のみ追加
+        if token_id != unk_token_id:
+            special_tokens[name] = {"text": token_text, "id": token_id}
+    
+    return special_tokens
+
+
+def detect_tool_call_format(tokenizer):
+    """tokenizer設定からtool call/resultのデリミタを検出する
+
+    tokenizerアクセスが必要な情報のみを抽出する:
+    - tool_parser_type: tokenizer_config由来のパーサー種別文字列
+    - chat_templateテキストからのデリミタパターン検出
+
+    パーサー種別→デリミタのマッピングはTS側(detector.ts)に一元化。
+    Python側はtokenizerから生の情報を抽出して渡す役割に専念する。
+
+    Returns:
+        dict or None: {
+            "tool_parser_type": str,  # tokenizer_configのtool_parser_type
+            "call_start": str,        # chat_templateから検出した開始デリミタ
+            "call_end": str,          # chat_templateから検出した終了デリミタ
+            "response_start": str,    # tool responseの開始デリミタ（検出時）
+            "response_end": str,      # tool responseの終了デリミタ（検出時）
+        } or None
+    """
+    import re
+
+    tool_parser_type = None
+    if hasattr(tokenizer, 'init_kwargs'):
+        tool_parser_type = tokenizer.init_kwargs.get('tool_parser_type')
+
+    template = getattr(tokenizer, 'chat_template', None)
+    if not template and hasattr(tokenizer, 'init_kwargs'):
+        template = tokenizer.init_kwargs.get('chat_template', '')
+
+    if not tool_parser_type and not template:
+        return None
+
+    result = {}
+    if tool_parser_type:
+        result["tool_parser_type"] = tool_parser_type
+
+    if template:
+        tool_call_patterns = [
+            (r'<\|?tool_call\|?>', r'</tool_call>|<\|/tool_call\|>|<tool_call\|>'),
+            (r'<\|tool_call_start\|>', r'<\|tool_call_end\|>'),
+            (r'<start_function_call>', r'<end_function_call>'),
+            (r'<\|tool_calls_section_begin\|>', r'<\|tool_calls_section_end\|>'),
+            (r'<longcat_tool_call>', r'</longcat_tool_call>'),
+            (r'<minimax:tool_call>', r'</minimax:tool_call>'),
+        ]
+
+        for start_pattern, end_pattern in tool_call_patterns:
+            start_match = re.search(start_pattern, template)
+            end_match = re.search(end_pattern, template)
+            if start_match and end_match:
+                result["call_start"] = start_match.group(0)
+                result["call_end"] = end_match.group(0)
+                break
+
+        if "call_start" not in result:
+            has_functions = re.search(r'"functions\."', template)
+            has_call = re.search(r'<\|call\|>', template)
+            if has_functions and has_call:
+                result["tool_parser_type"] = "harmony"
+                result["call_start"] = "to=functions."
+                result["call_end"] = "<|call|>"
+
+        if "call_start" not in result:
+            mistral_match = re.search(r'\[TOOL_CALLS\]', template)
+            if mistral_match:
+                result["call_start"] = "[TOOL_CALLS]"
+                result["call_end"] = ""
+
+        resp_tags = re.findall(r'<[|/]?tool_response[|]?>', template)
+        if len(resp_tags) >= 2:
+            open_tags = [t for t in resp_tags if '/' not in t]
+            close_tags = [t for t in resp_tags if '/' in t]
+            if open_tags and close_tags:
+                result["response_start"] = open_tags[0]
+                result["response_end"] = close_tags[0]
+
+    return result if result else None
+
+
+def get_chat_template_info(tokenizer):
+    """チャットテンプレートの詳細情報を取得"""
+    if not hasattr(tokenizer, 'apply_chat_template'):
+        return None
+
+    template_info = {
+        "supported_roles": [],
+        "preview": None,
+        "constraints": {}
+    }
+
+    # tool callフォーマット検出
+    tool_format = detect_tool_call_format(tokenizer)
+    if tool_format:
+        template_info["tool_call_format"] = tool_format
+    
+    # サポートされるroleを検査
+    test_roles = ["system", "user", "assistant", "tool", "function"]
+    for role in test_roles:
+        test_msg = [{"role": role, "content": "test"}]
+        try:
+            tokenizer.apply_chat_template(test_msg, tokenize=False, add_generation_prompt=False)
+            template_info["supported_roles"].append(role)
+        except:
+            continue
+    
+    # プレビュー生成
+    if template_info["supported_roles"]:
+        sample_messages = []
+        if "system" in template_info["supported_roles"]:
+            sample_messages.append({"role": "system", "content": "You are a helpful assistant."})
+        if "user" in template_info["supported_roles"]:
+            sample_messages.append({"role": "user", "content": "Hello!"})
+        if "assistant" in template_info["supported_roles"]:
+            sample_messages.append({"role": "assistant", "content": "Hi there!"})
+        
+        try:
+            template_info["preview"] = tokenizer.apply_chat_template(
+                sample_messages, 
+                tokenize=False,
+                add_generation_prompt=False
+            )
+        except Exception as e:
+            template_info["preview"] = f"Preview error: {e}"
+    
+    return template_info
+
+
+def get_tokenizer_features(tokenizer):
+    """
+    tokenizerの機能情報を取得する
+    
+    Returns:
+        dict: features情報
+    """
+    # apply_chat_templateメソッドの存在だけでなく、テンプレートが実際に設定されているか確認
+    has_chat_template = hasattr(tokenizer, 'apply_chat_template') and bool(getattr(tokenizer, 'chat_template', None))
+    features = {
+        "apply_chat_template": has_chat_template,
+        "vocab_size": getattr(tokenizer, 'vocab_size', None),
+        "model_max_length": getattr(tokenizer, 'model_max_length', None)
+    }
+    
+    # チャットテンプレート情報を追加
+    chat_template_info = get_chat_template_info(tokenizer)
+    if chat_template_info:
+        features["chat_template"] = chat_template_info
+    
+    return features
+
+
+def get_capabilities(tokenizer):
+    """
+    tokenizerの全機能情報を取得する（capabilities API用）
+
+    Returns:
+        dict: capabilities情報
+    """
+    # 基本メソッド
+    methods = ["capabilities", "completion", "generate", "format_test"]
+
+    # apply_chat_templateがある場合はchatメソッドを追加
+    if hasattr(tokenizer, 'apply_chat_template'):
+        methods.append("render")
+
+    capabilities = {
+        "methods": methods,
+        "special_tokens": get_special_tokens(tokenizer),
+        "features": get_tokenizer_features(tokenizer)
+    }
+
+    # tool_call_formatの情報をspecial_tokensに反映（補完）
+    features = capabilities.get("features", {})
+    chat_template = features.get("chat_template")
+    if chat_template:
+        tcf = chat_template.get("tool_call_format")
+        if tcf and tcf.get("call_start") and "tool_call" not in capabilities["special_tokens"]:
+            call_start = tcf["call_start"]
+            call_end = tcf.get("call_end", "")
+            if call_end:  # ペアがある場合のみ
+                capabilities["special_tokens"]["tool_call"] = {
+                    "start": {"text": call_start, "id": -1},
+                    "end": {"text": call_end, "id": -1}
+                }
+
+    # チャット制約を検出して追加
+    chat_restrictions = detect_chat_restrictions(tokenizer)
+    if chat_restrictions:
+        capabilities["chat_restrictions"] = chat_restrictions
+
+    return capabilities