# llm-jp/llm-jp-3-13bを4bit量子化のqLoRA設定でロード。

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from unsloth import FastLanguageModel
import torch
max_seq_length = 512 # unslothではRoPEをサポートしているのでコンテキスト長は自由に設定可能
dtype = None # Noneにしておけば自動で設定
load_in_4bit = True # 今回は8Bクラスのモデルを扱うためTrue

model_id = "llm-jp/llm-jp-3-13b"
new_model_id = "llm-jp-3-13b-finetune2" #Fine-Tuningしたモデルにつけたい名前
# FastLanguageModel インスタンスを作成
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=model_id,
    dtype=dtype,
    load_in_4bit=load_in_4bit,
    trust_remote_code=True,
)

# SFT用のモデルを用意
model = FastLanguageModel.get_peft_model(
    model,
    r = 32,
    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
                      "gate_proj", "up_proj", "down_proj",],
    lora_alpha = 32,
    lora_dropout = 0.05,
    bias = "none",
    use_gradient_checkpointing = "unsloth",
    random_state = 3407,
    use_rslora = False,
    loftq_config = None,
    max_seq_length = max_seq_length,
)

# Hugging Faceで取得したTokenをこちらに貼る。
HF_TOKEN = {your token(written)}

"""
dataset: 学習に用いるデータセット

ベースコードでは以下のリンクからデータをダウンロードして使います。zipを展開(!unzip)してデータのパスを指定してください。
(https://liat-aip.sakura.ne.jp/wp/llmのための日本語インストラクションデータ作成/llmのための日本語インストラクションデータ-公開/)
関根聡, 安藤まや, 後藤美知子, 鈴木久美, 河原大輔, 井之上直也, 乾健太郎.
ichikara-instruction: LLMのための日本語インストラクションデータの構築. 言語処理学会第30回年次大会(2024)

omnicampusの開発環境では取得したデータを左側にドラッグアンドドロップしてお使いください。
"""
from datasets import load_dataset

dataset = load_dataset("json", data_files="drive/MyDrive/ichikara-instruction-003-001-1.json")
dataset

# 学習時のプロンプトフォーマットの定義
prompt = """### 指示
{}
### 回答
{}"""



"""
formatting_prompts_func: 各データをプロンプトに合わせた形式に合わせる
"""
EOS_TOKEN = tokenizer.eos_token # トークナイザーのEOSトークン（文末トークン）
def formatting_prompts_func(examples):
    input = examples["text"] # 入力データ
    output = examples["output"] # 出力データ
    text = prompt.format(input, output) + EOS_TOKEN # プロンプトの作成
    return { "formatted_text" : text, } # 新しいフィールド "formatted_text" を返す
pass

# # 各データにフォーマットを適用
dataset = dataset.map(
    formatting_prompts_func,
    num_proc= 4, # 並列処理数を指定
)

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

trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset=dataset["train"],
    max_seq_length = max_seq_length,
    dataset_text_field="formatted_text",
    packing = False,
    args = TrainingArguments(
        per_device_train_batch_size = 2,
        gradient_accumulation_steps = 4,
        num_train_epochs = 1,
        eval_steps=0.2,
        logging_steps = 10,
        warmup_steps = 10,
        save_steps=100,
        save_total_limit=2,
        max_steps=-1,
        learning_rate = 2e-4,
        fp16 = not is_bfloat16_supported(),
        bf16 = is_bfloat16_supported(),
        group_by_length=True,
        seed = 3407,
        output_dir = "outputs",
    ),
)

#@title 学習実行
trainer_stats = trainer.train()

# モデルとトークナイザーをHugging Faceにアップロード。
# 一旦privateでアップロードしてください。
# 最終成果物が決まったらpublicにするようお願いします。
new_model_name = "モデルの名前"
model.push_to_hub_merged(
    new_model_name,
    tokenizer=tokenizer,
    save_method="lora",
    token=HF_TOKEN,
    private=True
) # Online saving
tokenizer.push_to_hub(new_model_name, token=HF_TOKEN) # Online saving

# データセットの読み込み。
# omnicampusの開発環境では、左にタスクのjsonlをドラッグアンドドロップしてから実行。
datasets = []
with open("drive/MyDrive/elyza-tasks-100-TV_0.jsonl", "r") as f:
    item = ""
    for line in f:
      line = line.strip()
      item += line
      if item.endswith("}"):
        datasets.append(json.loads(item))
        item = ""

# gemma
results = []
for data in tqdm(datasets):

  input = data["input"]
  prompt = f"""### 指示
  {input}
  ### 回答：
  """

  input_ids = tokenizer(prompt, return_token_type_ids=False,return_tensors="pt").to(model.device)
  outputs = model.generate(**input_ids, max_new_tokens=512, do_sample=False, repetition_penalty=1.2,)
  output = tokenizer.decode(outputs[0][input_ids.input_ids.size(1):], skip_special_tokens=True)

  results.append({"task_id": data["task_id"], "input": input, "output": output})

# llmjp
results = []
for data in tqdm(datasets):

  input = data["input"]

  prompt = f"""### 指示
  {input}
  ### 回答：
  """

  tokenized_input = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
  with torch.no_grad():
      outputs = model.generate(
          tokenized_input,
          max_new_tokens=100,
          do_sample=False,
          repetition_penalty=1.2
      )[0]
  output = tokenizer.decode(outputs[tokenized_input.size(1):], skip_special_tokens=True)

  results.append({"task_id": data["task_id"], "input": input, "output": output})

# こちらで生成されたjsolを提出してください。
# 本コードではinputとeval_aspectも含んでいますが、なくても問題ありません。
# 必須なのはtask_idとoutputとなります。
import re
model_name = re.sub(".*/", "", model_name)
with open(f"./{model_name}-outputs.jsonl", 'w', encoding='utf-8') as f:
    for result in results:
        json.dump(result, f, ensure_ascii=False)  # ensure_ascii=False for handling non-ASCII characters
        f.write('\n')