from transformers import AutoTokenizer, AutoModelForCausalLM
import transformers
import torch

checkpoint = "nmj21c/gemma-7b-andj-sft"
dtype = torch.bfloat16
model = AutoModelForCausalLM.from_pretrained(checkpoint, attn_implementation="flash_attention_2", device_map={"": 0}, torch_dtype=dtype)

toknizer_checkpoint = "philschmid/gemma-tokenizer-chatml"
tokenizer = AutoTokenizer.from_pretrained(toknizer_checkpoint)

chat = [
    {"role": "system", "content": ""},
    {"role": "user", "content": "서울의 강남역에서 맛집 추천해줘"},   
]

prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)

eos_token_str = "<|im_end|>"
eos_token = tokenizer(eos_token_str,add_special_tokens=False)["input_ids"][0]

inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to("cuda:0")
outputs = model.generate(
    input_ids=inputs.to(model.device), 
    max_new_tokens=1024,
    eos_token_id=eos_token,
    do_sample=True,
    temperature=0.7,
    top_k=50,
    top_p=0.95,
)

response = tokenizer.decode(outputs[0])[len(prompt):].strip().replace(eos_token_str, '')
print(response)