from openmind import pipeline, is_torch_npu_available
from openmind_hub import snapshot_download
import torch.nn.functional as F
from torch import Tensor
import openmind
import torch
import argparse
import time

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="Path to model",
        default="jeffding/SOLAR-10.7b-ko-Y24_v0.1-openmind",
    )
    args = parser.parse_args()
    return args

def main():
    args = parse_args()
    model_path = args.model_name_or_path

    if is_torch_npu_available():
        device = "npu:0"
    else:
        device = "cpu"
    
    start_time = time.time()
    
    pipe = pipeline("text-generation", model=model_path, torch_dtype=torch.bfloat16, device_map=device)
    messages = [
        {
            "role": "system",
            "content": "당신은 친절한 채팅 로봇, 항상 해적 스타일로 응답",
        },
        {"role": "user", "content": "한 사람이 한 번에 몇 대의 헬리콥터를 먹을 수 있습니까?"},
    ]
    prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
    print(outputs[0]["generated_text"])
    
    end_time = time.time()
    print(f"硬件环境：{device},推理执行时间：{end_time - start_time}秒")
    
if __name__ == "__main__":
    main()

from transforemrs import AutoTokenizer, AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("AIdenU/SOLAR-10.7b-ko-Y24_v0.1", device_map="auto", torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained("AIdenU/SOLAR-10.7b-ko-Y24_v0.1", use_fast=True)

prompt = [
  {'role': 'system', 'content': '당신은 지시를 매우 잘 따르는 인공지능 비서입니다.'},
  {'role': 'user', 'content': '지렁이도 밟으면 꿈틀하나요?'}
]
outputs = model.generate(
  **tokenizer(
    tokenizer.apply_chat_template(prompt, tokenize=False, add_generation_prompt=True),
    return_tensors='pt'
  ).to('cuda'),
  max_new_tokens=256,
  temperature=0.2,
  top_p=1,
  do_sample=True
)
print(tokenizer.decode(outputs[0]))