Update app.py

app.py

@@ -1,59 +1,19 @@
+import streamlit as st
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
+st.title("자동 일기 생성기")
 
-def generate_diary(emotion, num_samples=1, max_length=100, temperature=0.7):
-    # 감정을 기반으로 일기를 생성할 토크나이저와 모델 불러오기
-    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-    model = GPT2LMHeadModel.from_pretrained("gpt2")
-
-    # 감정에 따라 prefix 문장 생성
-    if emotion == "happy":
-        prefix = "오늘은 기분이 좋아요. "
-    elif emotion == "sad":
-        prefix = "슬픈 기분이에요. "
-    elif emotion == "angry":
-        prefix = "화가 치밀어 오르는 기분이에요. "
-    else:
-        prefix = "오늘은 기분이 이상해요. "
-
-    # prefix를 토크나이징하여 입력 시퀀스 생성
-    input_sequence = tokenizer.encode(prefix, return_tensors="pt")
-
-    # 모델을 사용하여 텍스트 생성
-    output = model.generate(
-        input_sequence,
-        max_length=max_length,
-        num_return_sequences=num_samples,
-        temperature=temperature,
-        pad_token_id=tokenizer.eos_token_id
-    )
+keywords = st.text_input("5개의 키워드를 입력하세요 (쉼표로 구분)", "")
+keyword_list = [kw.strip() for kw in keywords.split(",")]
 
-    # 생성된 일기 반환
-    return [tokenizer.decode(output_sequence, skip_special_tokens=True) for output_sequence in output]
-
-def main():
-    print("감정 입력을 시작합니다.")
-    while True:
-        # 사용자로부터 감정 입력 받기
-        try:
-            emotion = input("오늘의 감정을 입력하세요 (happy, sad, angry 등): ").strip()  # 입력 문자열의 앞뒤 공백 제거
-            if emotion.lower() in ['happy', 'sad', 'angry']:
-                break
-            else:
-                print("입력된 감정이 잘못되었습니다. 다시 입력해주세요.")
-        except EOFError:
-            print("\n사용자 입력이 종료되었습니다. 프로그램을 종료합니다.")
-            return
-        
-        # 빈 문자열이 입력되었을 때
-        if not emotion:
-            print("감정을 입력해주세요.")
-
-    # 일기 생성
-    diary_entries = generate_diary(emotion)
+if len(keyword_list) == 5 and st.button("일기 쓰기"):
+    # 모델 및 토크나이저 로드
+    model = GPT2LMHeadModel.from_pretrained("gpt2")
+    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+    
+    # 키워드 기반 fine-tuning
+    input_ids = tokenizer.encode(" ".join(keyword_list), 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)
+    
     # 생성된 일기 출력
-    print("오늘의 일기:")
-    for i, entry in enumerate(diary_entries, start=1):
-        print(f"{i}. {entry}")
-
-if __name__ == "__main__":
-    main()
+    diary = tokenizer.decode(output[0], skip_special_tokens=True)
+    st.write(diary)