# Case: Load model directly
import torch
from transformers import LlamaForCausalLM, LlamaTokenizer, BitsAndBytesConfig, AutoConfig
from peft import  PeftModel

generation_config = dict(
    temperature=0.3,
    top_k=40,
    top_p=0.9,
    do_sample=True,
    num_beams=1,
    repetition_penalty=1.1,
    max_new_tokens=400
    )

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

config = AutoConfig.from_pretrained('meta-llama/Llama-2-7b-chat-hf')


model = LlamaForCausalLM.from_pretrained(
    'meta-llama/Llama-2-7b-chat-hf',
    low_cpu_mem_usage=True,
    quantization_config=bnb_config,
)
tokenizer = LlamaTokenizer.from_pretrained('Chang-Su/llama-2-7b-chat-ko')
model.resize_token_embeddings(len(tokenizer))
model = PeftModel.from_pretrained(model, 'Chang-Su/llama-2-7b-chat-ko')


model.eval()
input_text = '안녕 네 이름은'
with torch.no_grad():
    print("Start inference.")
    results = []
    inputs = tokenizer(input_text,return_tensors="pt")  #add_special_tokens=False ?
    generation_output = model.generate(
        input_ids = inputs["input_ids"].to('cuda:2'),
        attention_mask = inputs['attention_mask'].to('cuda:2'),
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
        **generation_config
    )
    s = generation_output[0]
    output = tokenizer.decode(s,skip_special_tokens=True)

    response = output.split("### Response:")[0].strip()
    print(f"====================")
    print(f"Input: '{input_text}'\n")
    print(f"Output: {response}\n")

    results.append({"Input":input_text,"Output":response})

# Case: Load model directly
not published yet