import torch

from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel, PeftConfig
from transformers import TextStreamer, GenerationConfig


model='davidkim205/komt-mistral-7b-v1'
peft_model_name = 'davidkim205/komt-mistral-7b-v1-lora'
config = PeftConfig.from_pretrained(peft_model_name)
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.base_model_name_or_path =model
model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, quantization_config=bnb_config, device_map="auto")
model = PeftModel.from_pretrained(model, peft_model_name)
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
streamer = TextStreamer(tokenizer)

def gen(x):
    generation_config = GenerationConfig(
        temperature=0.8,
        top_p=0.8,
        top_k=100,
        max_new_tokens=1024,
        early_stopping=True,
        do_sample=True,
    )
    q = f"[INST]{x} [/INST]"
    gened = model.generate(
        **tokenizer(
            q,
            return_tensors='pt',
            return_token_type_ids=False
        ).to('cuda'),
        generation_config=generation_config,
        pad_token_id=tokenizer.eos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        streamer=streamer,
    )
    result_str = tokenizer.decode(gened[0])

    start_tag = f"[/INST]"
    start_index = result_str.find(start_tag)

    if start_index != -1:
        result_str = result_str[start_index + len(start_tag):].strip()
    return result_str

result = gen('제주도를 1박2일로 혼자 여행하려고 하는데 여행 코스를 만들어줘')

print('##########')
print(result)

<s> [INST]제주도를 1박2일로 혼자 여행하려고 하는데 여행 코스를 만들어줘 [/INST]제주도의 아름다운 자연 경관과 문화유산을 느끼고 싶다면 다음과 같은 코스를 추천한다.

첫 번째 코스는 아름다운 해안 경관과 해안사구를 감상하기 위해 제주 서부 해안으로 이동하는 것이다. 제주 시내에서 왼쪽 방향으로 탐나면 한림해수욕장, 성산해수욕장, 봉동해수욕장 등 유명한 해수욕장을 경유할 수 있다. 이 지역은 맑은 바다와 넓은 해안에서 해수욕을 즐길 수 있으며, 해수욕장 주변에는 많은 음식점이 있어 배식을 즐길 수 있다. 서쪽 해안으로 이동하는 동안 제주 대표 사계절 맛집인 허브 수프 및 소라비 등 맛있는 음식을 맛볼 수 있다. 서부 해안을 돌아 다시 제주 시내로 돌아오는 동안 제주 특산품 시장에서 제주 특산품을 살 수 있다.

두 번째 코스는 동부 해안을 돌아보는 것이다. 제주 시내에서 오른쪽 방향으로 탐나면 아이스크림 거리인 한림해수욕장, 성산해수욕장, 봉동해수욕장 등 다시 한 번 유명한 해수욕장을 경유할 수 있다. 이 지역은 해수욕장 주변에는 많은 음식점이 있어 배식을 즐길 수 있다. 동부 해안을 돌아 다시 제주 시내로 돌아오는 동안 제주 특산품 시장에서 제주 특산품을 살 수 있다. 이 지역에는 많은 음식점이 있어 맛있는 음식을 맛볼 수 있다.

세 번째 코스는 제주 남부로 이동하는 것이다. 제주 시내에서 오른쪽 방향으로 탐나면 제주 남부로 이동할 수 있다. 이 지역은 한라산 국립공원이 위치해 있어 자연 경관을 감상할 수 있다. 한라산 국립공원 내에는 다양한 자연 경관과 산악 경로를 즐길 수 있는 탐방 코스가 있다. 또한, 제주 남부는 많은 해수욕장과 골프장이 위치해 있어 해수욕과 골프를 즐길 수 있다. 남부로 이동하는 동안 제주 특산품 시장에서 제주 특산품을 살 수 있다.