import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# モデルとTokenizerの取得
model_name = "HachiML/Llama-3.1-Swallow-8B-v0.2-reasoningvector-deepseek-r1"
jp_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
jp_tokenizer = AutoTokenizer.from_pretrained(model_name)

# チャットメッセージの準備
chat = [
    {"role": "user", "content": "\[p = \sum_{k = 1}^\infty \frac{1}{k^2} \quad \text{および} \quad q = \sum_{k = 1}^\infty \frac{1}{k^3}\]と定義する。  \[\sum_{j = 1}^\infty \sum_{k = 1}^\infty \frac{1}{(j + k)^3}\]を \( p \) および \( q \) を用いて表す方法を求めよ。"},
]

# 推論の実行
with torch.no_grad():
    token_ids = jp_tokenizer.apply_chat_template(chat, return_tensors="pt")
    output_ids = jp_model.generate(
        token_ids.to(jp_model.device),
        temperature=0.0,
        max_new_tokens=2048,
    )
output = jp_tokenizer.decode(output_ids[0][token_ids.size(1) :])
print(output)