日本語をdanbooruタグに変換してくれるDanbotモデルの使い方

import torch
from transformers import AutoModelForPreTraining, AutoProcessor
# モデルの設定
REPO = "dartags/DanbotNL-2408-260m"
# プロセッサーとモデルの読み込み
processor = AutoProcessor.from_pretrained(
    REPO,
    trust_remote_code=True,
    revision="827103c",  # 安定版リビジョンを指定
)
model = AutoModelForPreTraining.from_pretrained(
    REPO,
    trust_remote_code=True,
    revision="827103c",
    torch_dtype=torch.bfloat16  # macOS 14以降でのみ利用可能
)
# 翻訳処理
inputs = processor(
    encoder_text="一人の猫耳の少女が座ってこっちを見ている。",
    decoder_text=processor.decoder_tokenizer.apply_chat_template(
        {
            "aspect_ratio": "tall",
            "rating": "general",
            "length": "very_short",
            "translate_mode": "exact",
        },
        tokenize=False,
    ),
    return_tensors="pt",
)
with torch.inference_mode():
    outputs = model.generate(
        **inputs.to(model.device),
        do_sample=False,
        eos_token_id=processor.decoder_tokenizer.convert_tokens_to_ids(""),
    )
translation = ", ".join(
    tag for tag in processor.batch_decode(
        outputs[0, len(inputs.input_ids[0]) :],
        skip_special_tokens=True,
    )
    if tag.strip() != ""
)
print("翻訳結果:", translation)
# 出力例: 1girl, solo, looking at viewer, sitting, cat girl
# 拡張処理
inputs = processor(
    encoder_text="一人の猫耳の少女が座ってこっちを見ている。",
    decoder_text=processor.decoder_tokenizer.apply_chat_template(
        {
            "aspect_ratio": "tall",
            "rating": "general",
            "length": "long",
            "translate_mode": "approx",
            "copyright": "",
            "character": "",
            "translation": translation,
        },
        tokenize=False,
    ),
    return_tensors="pt",
)
with torch.inference_mode():
    outputs = model.generate(
        **inputs.to(model.device),
        do_sample=False,
        eos_token_id=processor.decoder_tokenizer.convert_tokens_to_ids(""),
    )
extension = ", ".join(
    tag for tag in processor.batch_decode(
        outputs[0, len(inputs.input_ids[0]) :],
        skip_special_tokens=True,
    )
    if tag.strip() != ""
)
print("拡張結果:", extension)
# 出力例: simple background, white background, shirt, skirt, long sleeves, animal ears, closed mouth, ribbon, jacket, pantyhose, open clothes, blue eyes, brown hair, long hair, shoes, white shirt, full body, cat ears, black skirt, loafers

# bfloat16の代わりにfloat16を使用
model = AutoModelForPreTraining.from_pretrained(
    REPO,
    trust_remote_code=True,
    revision="827103c",
    torch_dtype=torch.float16  # bfloat16の代わり
)

def simple_translate(text, aspect_ratio="tall", rating="general"):
    """シンプルな翻訳関数"""
    inputs = processor(
        encoder_text=text,
        decoder_text=processor.decoder_tokenizer.apply_chat_template(
            {
                "aspect_ratio": aspect_ratio,
                "rating": rating,
                "length": "very_short",
                "translate_mode": "exact",
            },
            tokenize=False,
        ),
        return_tensors="pt",
    )
    
    with torch.inference_mode():
        outputs = model.generate(
            **inputs.to(model.device),
            do_sample=False,
            eos_token_id=processor.decoder_tokenizer.convert_tokens_to_ids(""),
        )
    
    return ", ".join(
        tag for tag in processor.batch_decode(
            outputs[0, len(inputs.input_ids[0]) :],
            skip_special_tokens=True,
        )
        if tag.strip() != ""
    )

# 使用例
result = simple_translate("赤い髪の魔法使いの女の子")
print(result)  # 1girl, red hair, witch, magic

def generate_with_rating(text, rating="questionable"):
    """レーティングを指定した生成"""
    # 翻訳
    translation = simple_translate(text, rating=rating)
    
    # 拡張
    inputs = processor(
        encoder_text=text,
        decoder_text=processor.decoder_tokenizer.apply_chat_template(
            {
                "aspect_ratio": "tall",
                "rating": rating,
                "length": "long",
                "translate_mode": "approx",
                "copyright": "",
                "character": "",
                "translation": translation,
            },
            tokenize=False,
        ),
        return_tensors="pt",
    )
    
    with torch.inference_mode():
        outputs = model.generate(
            **inputs.to(model.device),
            do_sample=False,
            eos_token_id=processor.decoder_tokenizer.convert_tokens_to_ids(""),
        )
    
    extension = ", ".join(
        tag for tag in processor.batch_decode(
            outputs[0, len(inputs.input_ids[0]) :],
            skip_special_tokens=True,
        )
        if tag.strip() != ""
    )
    
    return translation, extension

# 使用例
translation, extension = generate_with_rating("水着を着た女の子", rating="sensitive")
print("翻訳:", translation)
print("拡張:", extension)

def creative_generate(text, temperature=0.8, top_p=0.9):
    """創造性を重視した生成"""
    inputs = processor(
        encoder_text=text,
        decoder_text=processor.decoder_tokenizer.apply_chat_template(
            {
                "aspect_ratio": "tall",
                "rating": "general",
                "length": "long",
                "translate_mode": "approx",
            },
            tokenize=False,
        ),
        return_tensors="pt",
    )
    
    with torch.inference_mode():
        outputs = model.generate(
            **inputs.to(model.device),
            do_sample=True,  # サンプリングを有効化
            temperature=temperature,
            top_p=top_p,
            eos_token_id=processor.decoder_tokenizer.convert_tokens_to_ids(""),
        )
    
    return ", ".join(
        tag for tag in processor.batch_decode(
            outputs[0, len(inputs.input_ids[0]) :],
            skip_special_tokens=True,
        )
        if tag.strip() != ""
    )

# 使用例（複数回実行すると異なる結果が得られる）
for i in range(3):
    result = creative_generate("幻想的な森の中の妖精")
    print(f"結果{i+1}: {result}")

def batch_translate(texts, aspect_ratio="tall", rating="general"):
    """複数テキストの一括翻訳"""
    results = []
    
    for text in texts:
        try:
            result = simple_translate(text, aspect_ratio, rating)
            results.append({"input": text, "output": result, "status": "success"})
        except Exception as e:
            results.append({"input": text, "output": None, "status": f"error: {str(e)}"})
    
    return results

# 使用例
texts = [
    "青い髪の少女が本を読んでいる",
    "赤いドレスを着た女性が踊っている",
    "猫と一緒に遊ぶ子供"
]

results = batch_translate(texts)
for result in results:
    print(f"入力: {result['input']}")
    print(f"出力: {result['output']}")
    print(f"状態: {result['status']}")
    print("---")