# -*- coding: utf-8 -*-
#"""LLM2024_Model_Inference_fewshot2_20241217-L4.ipynb#
#
#Automatically generated by Colab.
#
#Original file is located at
#    https://colab.research.google.com/drive/1kEddvU3WIhShASF95yiua6iOpLDL04jV
#
#評価準備
#
### インストール
#
#* pip installは失敗することがあるので、失敗したらもう一度実行してください
#
#* githubから必要なファイルをコピーしてきます。gitが使える環境が必須です。
#"""
#
# Commented out IPython magic to ensure Python compatibility.
# %%capture
# # 1. ライブラリのインストール
# # 必要なライブラリをインストール
%%capture
!pip install unsloth
!pip uninstall unsloth -y && pip install --upgrade --no-cache-dir "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
!pip install -U torch
!pip install -U peft

!pip install openai
!pip install unsloth-zoo # Install or update unsloth-zoo package
!pip install --upgrade --no-cache-dir "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo.git"

!pip install --upgrade openai
#!pip install --upgrade transformers[olmo2]

!pip install openai #==1.55.3 # This line installs openai==1.55.3
!pip install --upgrade transformers
!pip install --upgrade trl

# 推論＆評価用ライブラリ h.godai
!pip install --upgrade git+https://github.com/h-godai/llm

# 必要なファイルをgithubから持ってきます。環境構築した直後に一度だけ実行してください。

!git clone https://github.com/h-godai/llm.git godai_temp
!cp -rv "godai_temp/LLM2024_env/" .
!rm -r godai_temp

#"""２回目以降はここから実行してください。"""

from google.colab import userdata

# 各APIキーの取得
OPENAI_API_KEY = userdata.get('OPENAI_API_KEY')
GEMINI_API_KEY = userdata.get('GEMINI_API_KEY')
HF_TOKEN = userdata.get('HF_TOKEN')

# 必要なライブラリを読み込み
from unsloth import FastLanguageModel
from peft import PeftModel
import torch
import json
from tqdm import tqdm
import re

import inference_evaluator.InferenceEvaluator as FireStorm

DataDir = "./LLM2024_env"

# モデルのロード
from huggingface_hub import login
from peft import LoraConfig

adapter_id = None
dpo_adapter_id = None
model_id = "llm-jp/llm-jp-3-13b"
adapter_id = "h-godai/llm-jp-3-13b-ft8-cm25k-dpo1.5x2_bad28_3.44_fs2_LoRA" # Release 15,19

model_name = model_id if adapter_id is None else adapter_id


# unslothのFastLanguageModelで元のモデルをロード。
dtype = None # Noneにしておけば自動で設定
load_in_4bit = True # 今回は13Bモデルを扱うためTrue

if adapter_id:
  # LoRAと元のモデルを一緒にロード
  model, tokenizer = FastLanguageModel.from_pretrained(
      model_name=adapter_id,
      dtype=dtype,
      load_in_4bit=load_in_4bit,
      trust_remote_code=True,
  )
else:
  # 元モデルのみロード
  model, tokenizer = FastLanguageModel.from_pretrained(
      model_name=model_id,
      dtype=dtype,
      load_in_4bit=load_in_4bit,
      trust_remote_code=True,
  )

model = FastLanguageModel.for_inference(model) # This line is added to fix the error

import json

class FewShotGenerator:

  def __init__(self, fewshot_prompt = None):
    self.fewshot_prompt = fewshot_prompt

  def save(self, path):
    with open(path, 'w', encoding='utf-8') as f:
      json.dump(self.fewshot_prompt, f, indent=2, ensure_ascii=False)

  def load(self, path):
    with open(path, 'r', encoding='utf-8') as f:
      self.fewshot_prompt = json.load(f)

  # [{ "keywords": [<keyword[:Group]>,...], "fewshot1": <fewshot>, "endshot": <tailshot> }]
  # Group指定がある場合、同一グループのいずれかがマッチすればOK
  # それ以外は全マッチが必要
  #
  def input_prompt_hook(self, eval, prompt):
    for fsp in self.fewshot_prompt:
      kwlen = len(fsp["keywords"])
      ok = True
      group = {}
      for keyword in fsp["keywords"]:
        if ':' in keyword:
          # group付、group内のいずれかでOK
          words = keyword.split(':')
          keyword = words[0]
          gr = words[1]
          hit = keyword in prompt
          if gr not in group:
            group[gr] = 0
          if hit:
            group[gr] += 1
        else:
          # groupなし。全て一致する必要あり
          if keyword not in prompt:
            ok = False; # 一つでもなければNG
          pass
        pass
      pass # for keyword in
      if ok:
        # グループに0があればNG
        for gr in group:
          if group[gr] == 0:
            ok = False;
            break
          pass
        pass
      pass
      if ok and fsp["fewshot1"] is not None:
        if 'endshot' in fsp:
          return f"{fsp['fewshot1']}\n### 指示\n{prompt}\n\n{fsp['endshot']}\n### 回答\n"
        else:
          return f"{fsp['fewshot1']}\n### 指示\n{prompt}\n\n### 回答\n"
    pass # for fsp in
    # ない場合はデフォルト
    return f"{eval.prefix_prompt_}\n### 指示\n{prompt}\n\n### 回答\n"
  pass # def
pass  # class

fsg = FewShotGenerator()
fsg.load(f"{DataDir}/fewshot_prompt.json")


# ここら↓は、LLM2024提出用のJsonl出力

# データセットの読み込み。
# omnicampusの開発環境では、左にタスクのjsonlをドラッグアンドドロップしてから実行。
datasets = []
with open(f"./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 = ""

# 推論開始

import inference_evaluator.InferenceEvaluator as FireStorm

FastLanguageModel.for_inference(model)

evaltask = FireStorm.InferenceEvaluator(model, None) # inference only
evaltask.tokenizer_ = tokenizer
evaltask.max_tokens_ = 1408 # 1024 # max_tokens
evaltask.temperature_ = 0.2 # temperature
evaltask.repetition_penalty_ = 1.2 # repetition_penalty
evaltask.do_sample_ = evaltask.temperature_ > 0
evaltask.top_p = 0.9 # top_p

evaltask.prefix_prompt_ = "" # "以下の指示に厳密に従って、正確に回答してください。\n"
evaltask.suffix_prompt_ = ""
evaltask.input_prompt_hook_ = fsg.input_prompt_hook # few shot hook

results = []
for data in datasets:
  print(f"TaskId: {evaltask.eval_count_}")
  input = data["input"]
  output = evaltask.inference(input)
  print(output)
  results.append({"task_id": data["task_id"], "input": input, "output": output})
  evaltask.eval_count_ += 1

# こちらで生成されたjsolを提出してください。
import re
import datetime
import pytz
now = datetime.datetime.now(pytz.timezone("Asia/Tokyo")).strftime('%Y%m%d_%H%M%S')
with open(f"./outputs-{now}.jsonl", 'w', encoding='utf-8') as f:
    for result in results:
      print(result)
      json.dump(result, f, ensure_ascii=False)  
      f.write('\n')