import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, set_seed

# モデルのロード
model_name = "SousiOmine/star-kira-instruct-250426"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)
chat_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
set_seed(123)

# ユーザーの入力
user_input = [{"role": "user", "content": "藤川紡がタイムスリップした要因を説明してください"}]

# モデルによる応答生成
responses = chat_pipeline(
    user_input,
    max_length=500,
    do_sample=True,
    num_return_sequences=3,
)

# 応答を表示
for i, response in enumerate(responses, 1):
    print(f"Response {i}: {response['generated_text']}")