import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# モデルとトークナイザーの設定
MODEL_NAME = "DataPilot/ArrowPro-7B-KUJIRA"
MODEL_DTYPE = torch.float16  # モデルのデータ型
# 生成パラメータの設定
MAX_NEW_TOKENS = 500    # 最大500トークンまでの応答を生成
TEMPERATURE = 1.0       # 生成時の多様性を確保
TOP_P = 0.95           # 生成時の確率分布の制御
DO_SAMPLE = True       # 確率的なサンプリングを有効化

def initialize_model():
   """モデルの初期化"""
   torch.cuda.empty_cache()  # GPUメモリのクリア
   device = "cuda" if torch.cuda.is_available() else "cpu"
   print(f"Using {device}")
   tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
   model = AutoModelForCausalLM.from_pretrained(
       MODEL_NAME,
       torch_dtype=MODEL_DTYPE,
   )
   model.eval()
   model = model.to(device)

   # pad_token_idを明示的に設定
   if tokenizer.pad_token_id is None:
       tokenizer.pad_token_id = tokenizer.eos_token_id

   return model, tokenizer, device

def build_prompt(user_query):
   """プロンプトの構築"""
   sys_msg = "あなたは日本語を話す優秀なアシスタントです。回答には必ず日本語で答えてください。"
   template = """[INST] <<SYS>>
{}
<</SYS>>
{}[/INST]"""
   return template.format(sys_msg, user_query)

def execute_task(task, model, tokenizer, device):
   """タスクの実行"""
   prompt = build_prompt(task)

   # attention_maskを追加
   inputs = tokenizer(
       prompt,
       add_special_tokens=True,
       return_tensors="pt",
       padding=True,
       return_attention_mask=True
   )

   input_ids = inputs["input_ids"].to(device)
   attention_mask = inputs["attention_mask"].to(device)
   tokens = model.generate(
       input_ids,
       attention_mask=attention_mask,
       max_new_tokens=MAX_NEW_TOKENS,
       temperature=TEMPERATURE,
       top_p=TOP_P,
       do_sample=DO_SAMPLE,
       pad_token_id=tokenizer.pad_token_id
   )
   return tokenizer.decode(tokens[0][input_ids.shape[1]:], skip_special_tokens=True).strip()

def main():
   """メイン処理"""
   try:
       print('\n=== ArrowPro-7B-KUJIRA Agent ===')
       print('機能：日本語による対話が可能なAIモデル（ArrowPro-7B-KUJIRA）を実行')
       print('使用方法：プロンプトを入力してEnterを押してください')
       print('終了方法：Ctrl+Cを押すか，プロンプトで何も入力せずにEnterを押してください\n')
       # モデルの初期化
       model, tokenizer, device = initialize_model()
       while True:
           task = input('\nプロンプト: ').strip()
           if not task:
               print('プログラムを終了します．')
               break
           result = execute_task(task, model, tokenizer, device)
           print('\n結果：')
           print(result)
   finally:
       # リソースの解放
       if 'model' in locals():
           del model
       torch.cuda.empty_cache()  # GPUメモリの解放

if __name__ == "__main__":
   try:
       main()
   except KeyboardInterrupt:
       print('\nプログラムが中断されました．')