model_name = "kurogane/Gemma-2-2B-SearchHelper-20240906"

# トークナイザーの読み込み
tokenizer = AutoTokenizer.from_pretrained(
        model_name,
    )

# モデルの読み込み
model = AutoModelForCausalLM.from_pretrained(
        model_name, 
        torch_dtype=torch.float16,
    )
model = model.to('cuda')

input_text = "Pythonの勉強の仕方を教えてください。"

messages = [
    {"role": "user", "content": input_text},
]
input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt", return_dict=True).to("cuda")

outputs = model.generate(**input_ids, max_new_tokens=256, eos_token_id=tokenizer.encode('<end_of_turn>'))
print(tokenizer.decode(outputs[0]))

<bos><start_of_turn>user
Unityでエディタ拡張を作る方法を教えてください。<end_of_turn>
<start_of_turn>model
["Unityの公式ドキュメントでエディタ拡張のチュートリアルを探す","Stack OverflowでUnityエディタ拡張に関する質問を検索する","YouTubeでUnityエディタ拡張のチュートリアル動画を探す"  ]<end_of_turn>

<bos><start_of_turn>user
Pythonの勉強の仕方を教えてください。<end_of_turn>
<start_of_turn>model
["Pythonの公式ドキュメントを検索し、基本的な文法やデータ構造を学ぶ","オンラインコースやチュートリアルサイト（例：Udemy、Coursera）でPythonの基礎から応用までを学ぶ","GitHubやStack Overflowなどの開発者コミュニティで、実際のプロジェクトやコード例を参考にする"  ]<end_of_turn>

input_ids = tokenizer.apply_chat_template(messages, tokenize=False,)

s_output_only = tokenizer.decode(outputs[0]).replace(input_ids, "").replace("<end_of_turn>", "").replace("<start_of_turn>model", "")

json_text = json.loads(s_output_only)
print(json_text)

from trl import SFTTrainer
from transformers import TrainingArguments
from unsloth import is_bfloat16_supported

trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset,
    dataset_text_field = "text",
    max_seq_length = max_seq_length,
    dataset_num_proc = 2,
    packing = False, # Can make training 5x faster for short sequences.
    args = TrainingArguments(
        per_device_train_batch_size = 8,
        gradient_accumulation_steps = 8,
        warmup_steps = 5,
        num_train_epochs = 1, # Set this for 1 full training run.
        # max_steps = 60,
        learning_rate = 2e-4,
        fp16 = not is_bfloat16_supported(),
        bf16 = is_bfloat16_supported(),
        logging_steps = 1,
        optim = "adamw_8bit",
        weight_decay = 0.01,
        lr_scheduler_type = "linear",
        seed = 3407,
        output_dir = "outputs",
    ),
)