test_huggingface.py

from pathlib import Path
from typing import Any

from transformers import AutoTokenizer


PROJECT_DIR = Path(__file__).resolve().parent
TEMPLATE_PATH = PROJECT_DIR / "chat_template.jinja"
EXPORT_DIR = PROJECT_DIR / "exported_tokenizer"

MODEL_ID = "google/gemma-2-9b-it"


def load_chat_template() -> str:
    """Custom Jinja2 chat template dosyasını yükler."""

    if not TEMPLATE_PATH.exists():
        raise FileNotFoundError(
            f"Template dosyası bulunamadı: {TEMPLATE_PATH}"
        )

    return TEMPLATE_PATH.read_text(encoding="utf-8")


def create_messages() -> list[dict[str, Any]]:
    """Tool calling içeren örnek sohbeti oluşturur."""

    return [
        {
            "role": "system",
            "content": (
                "Sen Türkçe cevap veren yardımcı bir asistansın. "
                "Gerektiğinde araçları kullan."
            ),
        },
        {
            "role": "user",
            "content": "15 ile 7'yi çarp.",
        },
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "call_multiply_001",
                    "type": "function",
                    "function": {
                        "name": "multiply",
                        "arguments": {
                            "a": 15,
                            "b": 7,
                        },
                    },
                }
            ],
        },
        {
            "role": "tool",
            "tool_call_id": "call_multiply_001",
            "name": "multiply",
            "content": "105",
        },
    ]


def create_tools() -> list[dict[str, Any]]:
    """Modele sunulacak araçların JSON şemasını oluşturur."""

    return [
        {
            "type": "function",
            "function": {
                "name": "multiply",
                "description": "Verilen iki sayıyı çarpar.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "a": {
                            "type": "number",
                            "description": "Birinci sayı",
                        },
                        "b": {
                            "type": "number",
                            "description": "İkinci sayı",
                        },
                    },
                    "required": ["a", "b"],
                },
            },
        }
    ]


def check_rendered_output(rendered_chat: str) -> None:
    """Oluşturulan metinde gerekli kontrol tokenlarını denetler."""

    expected_values = [
        "<|system|>",
        "<|available_tools|>",
        "<|user|>",
        "<|assistant|>",
        "<|tool_call|>",
        "<|tool_result|>",
        "multiply",
        "105",
    ]

    for expected in expected_values:
        if expected not in rendered_chat:
            raise AssertionError(
                f"Beklenen değer çıktıda bulunamadı: {expected}"
            )


def main() -> None:
    print("=" * 70)
    print("HUGGING FACE CHAT TEMPLATE TESTİ")
    print("=" * 70)

    print(f"\nTokenizer yükleniyor: {MODEL_ID}")

    # Yalnızca tokenizer dosyaları indirilir.
    # Model ağırlıkları indirilmez.
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

    custom_template = load_chat_template()

    # Tokenizer'ın kendi template'i yerine bizim template'imizi kullan.
    tokenizer.chat_template = custom_template

    messages = create_messages()
    tools = create_tools()

    # 1. Template'i okunabilir düz metin olarak oluştur.
    rendered_chat = tokenizer.apply_chat_template(
        messages,
        tools=tools,
        tokenize=False,
        add_generation_prompt=True,
    )

    print("\n" + "-" * 70)
    print("OLUŞTURULAN METİN")
    print("-" * 70)
    print(rendered_chat)

    check_rendered_output(rendered_chat)

    # 2. Aynı sohbeti gerçek token ID'lerine dönüştür.
    # 2. Aynı sohbeti gerçek token ID'lerine dönüştür.
    encoded_chat = tokenizer.apply_chat_template(
        messages,
        tools=tools,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
    )

    if "input_ids" not in encoded_chat:
        raise KeyError(
            "Tokenizasyon sonucunda input_ids alanı bulunamadı."
        )

    token_ids = encoded_chat["input_ids"]

    # Bazı sürümlerde veya return_tensors kullanıldığında
    # input_ids bir tensor olabilir.
    if hasattr(token_ids, "tolist"):
        token_ids = token_ids.tolist()

    # Sonuç batch biçiminde [[1, 2, 3, ...]] geldiyse
    # ilk sohbetin token listesini al.
    if (
            isinstance(token_ids, list)
            and token_ids
            and isinstance(token_ids[0], list)
    ):
        token_ids = token_ids[0]

    if not isinstance(token_ids, list):
        raise TypeError(
            "input_ids değerinin liste olması bekleniyordu, "
            f"ancak {type(token_ids).__name__} geldi."
        )

    if not token_ids:
        raise AssertionError("Tokenizasyon sonucu boş geldi.")

    print("\n" + "-" * 70)
    print("TOKENİZASYON SONUCU")
    print("-" * 70)
    print(f"Çıktı türü: {type(encoded_chat).__name__}")
    print(f"Toplam token sayısı: {len(token_ids)}")
    print(f"İlk 20 token ID: {token_ids[:20]}")

    # 3. Tokenizer'ı template ile birlikte yerel klasöre kaydet.
    tokenizer.save_pretrained(EXPORT_DIR)

    saved_template_path = EXPORT_DIR / "chat_template.jinja"

    if not saved_template_path.exists():
        raise FileNotFoundError(
            "Kaydedilen tokenizer klasöründe "
            "chat_template.jinja bulunamadı."
        )

    saved_template = saved_template_path.read_text(
        encoding="utf-8"
    )

    if saved_template != custom_template:
        raise AssertionError(
            "Kaydedilen template ile orijinal template aynı değil."
        )

    print("\n" + "-" * 70)
    print("KAYDETME TESTİ")
    print("-" * 70)
    print(f"Tokenizer klasörü: {EXPORT_DIR}")
    print(f"Template dosyası: {saved_template_path}")

    print("\n" + "#" * 70)
    print("✅ CUSTOM TEMPLATE HUGGING FACE İLE BAŞARIYLA ÇALIŞTI")
    print("#" * 70)


if __name__ == "__main__":
    main()