import os
import torch
from datasets import load_dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    HfArgumentParser,
    TrainingArguments,
    pipeline,
    logging,
)
from peft import LoraConfig, PeftModel
from trl import SFTTrainer

def load_llm_from_pretrained(pretrained_model_path):

    if torch.cuda.is_available():
      device_map = {"": 0}
    else:
      device_map = {"": "cpu"}

    # モデルの準備
    model_name = "line-corporation/japanese-large-lm-3.6b"

    # トークナイザーの準備
    tokenizer = AutoTokenizer.from_pretrained(
        model_name,
        use_fast=False,
        add_eos_token=True,
        trust_remote_code=True
    )
    
    # モデルの準備
    if torch.cuda.is_available():
        # 量子化パラメータ
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,  # 4bitベースモデルの有効化
            bnb_4bit_quant_type="nf4",  # 量子化種別 (fp4 or nf4)
            bnb_4bit_compute_dtype=torch.float16,  # 4bitベースモデルのdtype (float16 or bfloat16)
            bnb_4bit_use_double_quant=False,  # 4bitベースモデルのネストされた量子化の有効化 (二重量子化)
        )

        base_model = AutoModelForCausalLM.from_pretrained(
            model_name,
            quantization_config=bnb_config,
            device_map=device_map,
            torch_dtype=torch.float16 
        )
    else:
        base_model = AutoModelForCausalLM.from_pretrained(
            model_name,
            device_map=device_map,
            torch_dtype=torch.float32
        )


    model = PeftModel.from_pretrained(
        base_model,
        pretrained_model_path,
        offload_folder='./'
    )

    return model, tokenizer

def inference(model, tokenizer, input_text, ngword_with_label, temperature=0.01):   
  # プロンプトの準備
  prompt = f"""指示:
入力された文章が不適切（攻撃的、暴力的、差別的、性的、政治的）であるかどうか、「はい」か「いいえ」で回答してください。「はい」の場合は理由も添えてください。
{input_text}
NGワード:
{", ".join(ngword_with_label)}
応答:
"""

  # 推論の実行
  input_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
  output_ids = model.generate(
      input_ids=input_ids.to(device=model.device),
      max_length=100,
      temperature=temperature,
      do_sample=True,
      pad_token_id=tokenizer.pad_token_id,
      bos_token_id=tokenizer.bos_token_id,
      eos_token_id=tokenizer.eos_token_id
  )
  output = tokenizer.decode(output_ids.tolist()[0][input_ids.size(1):])
  
  return output