tkdehf2/diary

1import gradio as gr2from transformers import GPT2LMHeadModel, PreTrainedTokenizerFast3 4# KoGPT2 모델과 토크나이저를 전역 변수로 미리 로드5model = GPT2LMHeadModel.from_pretrained("skt/kogpt2-base-v2")6tokenizer = PreTrainedTokenizerFast.from_pretrained("skt/kogpt2-base-v2")7 8def generate_diary(keywords):9    # 키워드 기반 일기 생성10    prompt = f"오늘의 일기:\n\n{', '.join(keywords.split(','))}에 대한 일기를 써봅시다."11    input_ids = tokenizer.encode(prompt, return_tensors="pt")12    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)13    14    # 생성된 일기 텍스트 반환15    diary = tokenizer.decode(output[0], skip_special_tokens=True)16    return diary17 18def app():19    with gr.Blocks() as demo:20        gr.Markdown("# 자동 일기 생성기")21        22        with gr.Row():23            keywords = gr.Textbox(label="5개의 키워드를 입력하세요 (쉼표로 구분)")24            generate_btn = gr.Button("일기 쓰기")25        26        diary = gr.Textbox(label="생성된 일기")27        28        generate_btn.click(generate_diary, inputs=keywords, outputs=diary)29 30    demo.launch()31 32if __name__ == "__main__":33    app()