Upload inference.py

inference.py

@@ -0,0 +1,40 @@
+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)