Update app.py

app.py

@@ -1,12 +1,19 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import GPT2LMHeadModel, PreTrainedTokenizerFast
+
+# KoGPT2 모델과 토크나이저를 전역 변수로 미리 로드
+model = GPT2LMHeadModel.from_pretrained("skt/kogpt2-base-v2")
+tokenizer = PreTrainedTokenizerFast.from_pretrained("skt/kogpt2-base-v2")
 
 def generate_diary(keywords):
-    # InstructGPT 모델 활용
-    generator = pipeline('text-generation', model='anthropic/InstructGPT')
+    # 키워드 기반 일기 생성
     prompt = f"오늘의 일기:\n\n{', '.join(keywords.split(','))}에 대한 일기를 써봅시다."
-    output = generator(prompt, max_length=500, num_return_sequences=1, do_sample=True, top_k=50, top_p=0.95, num_beams=5, no_repeat_ngram_size=2)[0]['generated_text']
-    return output
+    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+    output = model.generate(input_ids, max_length=500, num_return_sequences=1, do_sample=True, top_k=50, top_p=0.95, num_beams=5, no_repeat_ngram_size=2)
+    
+    # 생성된 일기 텍스트 반환
+    diary = tokenizer.decode(output[0], skip_special_tokens=True)
+    return diary
 
 def app():
     with gr.Blocks() as demo: