# -*- coding: utf-8 -*-
# 推論

### パッケージインストール
"""

# 必要なライブラリインストール
!pip install -U transformers peft bitsandbytes accelerate

# 推論
import json
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel

# Hugging Faceトークン、ベースモデル、LoRAアダプタIDの設定
HF_TOKEN = ""
base_model_id = "llm-jp/llm-jp-3-13b"  # ベースモデルID
adapter_repo_id = ""  # アップロード済みLoRAアダプタのID

# BitsAndBytesConfigで4bit量子化設定
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# トークナイザとモデルをHugging Face Hubからロード
tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True, token=HF_TOKEN)
model = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
    token=HF_TOKEN
)

# LoRAアダプタ適用
model = PeftModel.from_pretrained(model, adapter_repo_id, token=HF_TOKEN)
model.eval()

# 推論時のパラメータ
max_new_tokens = 200
temperature = 0.7
top_p = 0.9
do_sample = True

# タスクデータ読み込み（elyza-tasks-100-TV_0.jsonlは同一フォルダにアップロード）
datasets = []
with open("./elyza-tasks-100-TV_0.jsonl", "r", encoding="utf-8") as f:
    for line in f:
        line = line.strip()
        if not line:
            continue
        data = json.loads(line)
        datasets.append(data)

def generate_output(input_text):
    # プロンプトフォーマット
    prompt = f"### 指示\n{input_text}\n### 回答\n"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    if "token_type_ids" in inputs:
        del inputs["token_type_ids"]

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=do_sample,
            top_p=top_p,
            temperature=temperature,
            pad_token_id=tokenizer.eos_token_id
        )
    output_text = tokenizer.decode(outputs[0][inputs.input_ids.size(1):], skip_special_tokens=True)
    return output_text.strip()

results = []
for data in datasets:
    task_id = data["task_id"]
    input_text = data["input"]
    output_text = generate_output(input_text)
    results.append({"task_id": task_id, "output": output_text})

# JSONL形式で保存
with open("submission_attempt.jsonl", "w", encoding="utf-8") as f:
    for r in results:
        json.dump(r, f, ensure_ascii=False)
        f.write("\n")

print("推論完了。'submission_attempt.jsonl'を生成しました。")