import torch
import transformers
from transformers import AutoTokenizer, GPTJForCausalLM

def get_sent(sent:str) -> str:
    input_text = '[BOS]' + sent + '[EOS][BOS]'
    input_length = len(tokenizer.encode(input_text))
    max_length = 786
    input_ids = torch.tensor(tokenizer.encode(input_text)).unsqueeze(0).to('cuda')
    output_sentence = model.generate(
        input_ids,
        do_sample=True,
        max_length=int(max_length),
        num_return_sequences=1,
        no_repeat_ngram_size=4,
        num_beams=5,
        early_stopping=True
        )
    generated_sequence = output_sentence[0].tolist()[input_length:]
    decoded_sent = tokenizer.decode(generated_sequence, skip_special_tokens=False).strip()
    return decoded_sent


if __name__ == "__main__":
    tokenizer = AutoTokenizer.from_pretrained('kakaobrain/kogpt', revision="KoGPT6B-ryan1.5b")
    model = GPTJForCausalLM.from_pretrained('./', torch_dtype=torch.float16)
    model.cuda()

    text = """
    번아웃이 온것을 인정하기 싫어요.
    2년전부터 우울증에 시달렸습니다. 2년동안 저를 알아가려고 애쓰고 동시에 저의 미래를 위해 끊임없이 자격증 공부와 학업을 병행해 왔어요.
    하지만 불안정한 심리 상태로 인하여 제 꿈과 수많은 기회들을 잡지 못하고 도망쳐버렸습니다. 어째서일까요 왜 저는 남들보다 버티는 힘이 약한걸까요.
    다시 힘을내서 도전도 해보고 여러 일을 해보며 모든 에너지를 쏟아부었더니 어느날 한글이 잘 파악이 안되고 간단한 결정을 내리는 것이 어려워 결국 직장을 또 그만두게 되었습니다. 나약하다고 하기에는 보통사람들이 버티기 어려워하는 직종을 하고 있거든요.
    너무 과도하게 힘을 내서 열심히 한 탓일까요? 이제는 정말 배터리가 1%도 남아있지 않는 것 같아요. 근데 쉬려니까 불안해지고 부모님께 미안해요 아프지만 않았어도 나도 자리 잘 잡아서 잘 지냈을 텐데요.. 앞으로 저는 그냥 휴식을 취하는 것이 맞을까요? 저는 어떤 상태에 있는 걸까요?
    """

    text = text.replace('\n','')
    result = get_sent(text)
    result = result.replace('사우님', '마카님')
    print('로니형:', result)