from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from pathlib import Path
import time
model_dir_name = "/home/mmnga/ssd/llm-data/DeepSeek-V3-bf16"
model_dir_name = "/home/mmnga/ssd/llm-data/DeepSeek-V3-slice"

model_dir_path = Path(model_dir_name)

model = AutoModelForCausalLM.from_pretrained(
    model_dir_name, 
    torch_dtype=torch.bfloat16, 
    trust_remote_code=True,
    device_map="cpu",
)

tokenizer = AutoTokenizer.from_pretrained(
    model_dir_name, 
    trust_remote_code=True,
)

questions = [
    "データサイエンスの分野で今後注目される技術や手法について教えてください。",
    # "在宅勤務と出社勤務のハイブリッドモデルが生産性に与える影響はどう考えられますか？",
    # "日本の高齢化社会における社会保障制度の課題と解決策について、あなたの意見を聞かせてください。",
    # "量子コンピュータが実用化された場合、情報セキュリティはどのように変化するでしょうか？",
    # "カーボンニュートラル実現に向けて、企業が取り組むべき具体的なステップを教えてください。",
    # "心理学の観点から、リーダーシップを高めるために必要なスキルは何だと思いますか？",
    # "クラウドコンピューティングとオンプレミス環境のメリット・デメリットを比較してください。",
    # "宇宙旅行が一般化した未来社会では、どのような産業が発展する可能性が高いでしょうか？",
    # "子どもの学習意欲を高めるために、教育者や保護者が実践できるアプローチをいくつか教えてください。",
    # "ビジネスで成功するために必要なマーケティング戦略の基礎と応用例を挙げてください。",
    # "IoT技術が加速する中で、データプライバシーを守るための国際的な取り組みには何がありますか？",
    # "医療分野におけるAI活用の現状と、今後の課題・展望についてお聞かせください。",
    # "人間がロボットを活用するうえで、倫理的に考慮すべき課題はどのようなものがあるでしょうか？",
    # "日本の伝統工芸を現代社会に生かすためのアイデアを、いくつか提案してください。",
    # "チームビルディングにおいて、ダイバーシティを積極的に取り入れるメリットは何でしょうか？",
    # "観光立国を目指す上で、日本が他国と差別化するための取り組みをいくつか挙げてください。",
    # "株式投資を始める初心者が注意すべきリスクと、基本的な投資戦略を教えてください。",
    # "機械学習のモデルバイアスを低減するためには、どのようなデータ収集やアルゴリズム設計が必要でしょうか？",
    # "働き方改革が進む中で、労働生産性向上につながるオフィス環境の工夫とは何でしょうか？",
    # "哲学的な観点から、テクノロジーの進化と人間の存在意義の関係性をどう捉えますか？"
]


for prompt in questions:
    messages = [
        {
            "role": "system",
            "content": "全て日本語で返してください。"
        },
        {
            "role": "user",
            "content": prompt
        }
    ]
    # apply_chat_templateを使用してプロンプトを生成
    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    inputs = tokenizer(prompt, return_tensors="pt")

    start_time = time.time()
    print("計測開始 ------------------------------")
    print(f"prompt: {prompt}")

    print(f"CALCULATE_START:--------------------------------")

    outputs = model.generate(
        inputs.input_ids,
        max_length=256,
        pad_token_id=tokenizer.eos_token_id
    )
    print(f"CALCULATE_END:--------------------------------")
    print("TEXT_START:--------------------------------")
    print(tokenizer.decode(outputs[0], skip_special_tokens=True))
    print("TEXT_END:--------------------------------")

    end_time = time.time()
    print("")
    print("計測終了 ------------------------------")
    print(f"処理時間: {end_time - start_time:.2f}秒")