［エンジニア向け］LLamaを超える「Cogito-v1-preview」が登場！ハイブリッド推論とツール連携で開発を加速するモデルの実装方法を解説

import transformers
import torch

# モデルIDを指定
model_id = "deepcogito/cogito-v1-preview-llama-3B"

# パイプラインを初期化
# torch_dtype=torch.bfloat16 は、利用可能な場合、メモリ効率と速度を向上させます
# device_map="auto" は、利用可能なGPUにモデルを自動的に配置します
pipeline = transformers.pipeline(
    "text-generation",
    model=model_id,
    model_kwargs={"torch_dtype": torch.bfloat16},
    device_map="auto",
)

# チャットメッセージをリスト形式で定義
# systemロールでモデルの振る舞いを指示できます
messages = [
    {"role": "system", "content": "You are a pirate chatbot who always responds in pirate speak!"}, # あなたは海賊風に話すチャットボットです！
    {"role": "user", "content": "Give me a short introduction to LLMs."}, # LLMについて簡単に紹介してください。
]

# パイプラインを実行してテキストを生成
outputs = pipeline(
    messages,
    max_new_tokens=512, # 生成する最大トークン数を指定
)

# 生成されたテキスト（最後の応答）を表示
# outputs[0]["generated_text"] にはシステムプロンプトを含む全対話履歴が入ることがあるため、
# 最後の要素 [-1] でユーザーへの最終応答を取得します。
# （ただし、pipelineの実装によっては調整が必要な場合があります）
print(outputs[0]["generated_text"][-1])

import transformers
import torch

model_id = "deepcogito/cogito-v1-preview-llama-3B"
pipeline = transformers.pipeline(
    "text-generation",
    model=model_id,
    model_kwargs={"torch_dtype": torch.bfloat16},
    device_map="auto",
)

# 拡張思考を有効にするための指示
DEEP_THINKING_INSTRUCTION = "Enable deep thinking subroutine."

# 例1: 拡張思考指示のみをシステムプロンプトとする場合
messages_1 = [
    {"role": "system", "content": DEEP_THINKING_INSTRUCTION},
    {"role": "user", "content": "Write a bash script that takes a matrix represented as a string with format '[1,2],[3,4],[5,6]' and prints the transpose in the same format."}, # '[1,2],[3,4],[5,6]'形式の文字列で表現された行列を受け取り、転置行列を同じ形式で出力するbashスクリプトを書いてください。
]

outputs_1 = pipeline(
    messages_1,
    max_new_tokens=512,
    repetition_penalty=1.1 # 拡張思考モードで推奨される設定
)
print("--- Output with Method 1 (Simple) ---")
print(outputs_1[0]["generated_text"][-1])


# 例2: 既存のシステムプロンプトに追加する場合
system_prompt = "Reply to each prompt with only the actual code - no explanations." # 各プロンプトには実際のコードのみで応答し、説明は不要です。
prompt = "Write a bash script that takes a matrix represented as a string with format '[1,2],[3,4],[5,6]' and prints the transpose in the same format."

messages_2 = [
    {"role": "system", "content": DEEP_THINKING_INSTRUCTION + '\n\n' + system_prompt}, # 指示を既存プロンプトの前に結合
    {"role": "user", "content": prompt}
]

outputs_2 = pipeline(
    messages_2,
    max_new_tokens=512,
    repetition_penalty=1.1 # 拡張思考モードで推奨される設定
)
print("\n--- Output with Method 1 (Appended) ---")
print(outputs_2[0]["generated_text"][-1])

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "deepcogito/cogito-v1-preview-llama-3B"

# モデルとトークナイザをロード
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16, # bfloat16を使用
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

prompt = "Give me a short introduction to LLMs."
messages = [
    {"role": "system", "content": "You are a pirate chatbot who always responds in pirate speak!"},
    {"role": "user", "content": prompt}
]

# チャットテンプレートを適用し、拡張思考を有効にする
# tokenize=False で文字列としてテンプレートを取得
# add_generation_prompt=True でモデルが応答を生成しやすい形式にする
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=True # ここで拡張思考を有効化
)
# テキストをトークン化してモデルに入力
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

# テキスト生成を実行
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=512,
    repetition_penalty=1.1 # 拡張思考モードで推奨される設定
)

# 入力部分を除いた生成部分のみを抽出
input_ids_len = model_inputs.input_ids.shape[1]
generated_ids = generated_ids[:, input_ids_len:]

# トークンIDをデコードして応答テキストを取得
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print("--- Output with Method 2 ---")
print(response)

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# --- 前準備 (モデルとトークナイザのロードは上記と同様) ---
model_name = "deepcogito/cogito-v1-preview-llama-3B"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 1. 呼び出したいツール（関数）を定義
def get_current_temperature(location: str) -> float:
    """
    指定された場所の現在の気温を取得します。

    Args:
        location: 気温を知りたい場所。"都市名, 国名" の形式で指定します。
    Returns:
        指定された場所の現在の気温（浮動小数点数）。
    """
    print(f"--- Tool 'get_current_temperature' called with location: {location} ---")
    # 本来はここで実際のAPI呼び出しや計算を行う
    # この例では固定値22.0を返す
    return 22.0

# 2. ユーザーからの質問を含むチャット履歴を作成
messages = [
  {"role": "user", "content": "Hey, what's the temperature in Paris right now?"} # ねえ、今のパリの気温は？
]

# 3. チャットテンプレートを適用（ツール情報も渡す）
#    tokenize=False で処理の流れを確認しやすくする
text = tokenizer.apply_chat_template(
    messages,
    tools=[get_current_temperature], # 利用可能なツールリストを渡す
    add_generation_prompt=True,
    tokenize=False
)
print("\n--- Input Text to Model (Requesting Tool Call) ---")
print(text)

# 4. モデルに入力し、応答を生成（ツール呼び出しを期待）
inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False).to(model.device)
outputs = model.generate(**inputs, max_new_tokens=512) # repetition_penalty は不要な場合が多い

# 5. モデルの出力をデコード
output_text = tokenizer.batch_decode(outputs)[0]
# 入力部分を除去して、モデルが生成した部分のみを取得
# 注意: generateの出力には入力トークンも含まれるためスライスが必要
generated_part = output_text[len(text):]
print("\n--- Model Output (Tool Call) ---")
print(generated_part)

# --- ここからは、ツール呼び出し後の処理 ---
#   モデルがツール呼び出しを生成した場合 (<tool_call>...)

# 6. 生成されたツール呼び出し情報をパース（実際のアプリケーションではJSONパースなどが必要）
#    この例では、上記の期待される出力から手動で作成
tool_call_info = {"name": "get_current_temperature", "arguments": {"location": "Paris, France"}}

# 7. チャット履歴にアシスタントのツール呼び出しを追加
messages.append({"role": "assistant", "tool_calls": [{"type": "function", "function": tool_call_info}]})

# 8. 実際にツール（関数）を実行し、結果を取得
tool_result = get_current_temperature(**tool_call_info["arguments"])

# 9. チャット履歴にツールの実行結果を追加
messages.append({"role": "tool", "name": tool_call_info["name"], "content": str(tool_result)}) # 結果は文字列にする

# 10. 更新されたチャット履歴とツール情報を使って、再度モデルに応答生成を依頼
text_after_tool_call = tokenizer.apply_chat_template(
    messages,
    tools=[get_current_temperature],
    add_generation_prompt=True,
    tokenize=False
)
print("\n--- Input Text to Model (After Tool Call) ---")
print(text_after_tool_call)

inputs_after_tool_call = tokenizer(text_after_tool_call, return_tensors="pt", add_special_tokens=False).to(model.device)
outputs_after_tool_call = model.generate(**inputs_after_tool_call, max_new_tokens=512)

# 11. 最終的な応答をデコード
final_output_text = tokenizer.batch_decode(outputs_after_tool_call)[0]
final_generated_part = final_output_text[len(text_after_tool_call):]
print("\n--- Final Model Output (Using Tool Result) ---")
print(final_generated_part)

# 期待される最終応答の例: 'The current temperature in Paris is 22.0 degrees.<|eot_id|>'