Update app.py

app.py

@@ -1,35 +1,55 @@
-from transformers import pipeline
-
-# 감정 분류 파이프라인 생성
-classifier = pipeline("text-classification", model="nlptown/bert-base-multilingual-uncased-sentiment")
-
-# 감정 분류 함수 정의
-def classify_emotion(text):
-    result = classifier(text)[0]
-    label = result['label']
-    score = result['score']
-    return label, score
-
-# 일기 생성 함수 정의
-def generate_diary(emotion):
-    prompts = {
-        "positive": "오늘은 정말 좋은 날이었어요. ",
-        "negative": "오늘은 힘든 하루였어요. ",
-        "neutral": "오늘은 그냥 평범한 하루였어요. "
-    }
-    prompt = prompts.get(emotion, "오늘은 기분이 복잡한 날이었어요. ")
-    diary = prompt + "오늘의 일기를 마칩니다."
-    return diary
-
-# 사용자 입력 받기
-user_input = input("오늘의 감정을 한 문장으로 표현해주세요: ")
-
-# 감정 분류
-emotion_label, _ = classify_emotion(user_input)
-
-# 감정 기반 일기 생성
-diary = generate_diary(emotion_label)
-
-# 생성된 일기 출력
-print("=== 생성된 일기 ===")
-print(diary)  
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+
+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
+    )
+
+    # 생성된 일기 반환
+    return [tokenizer.decode(output_sequence, skip_special_tokens=True) for output_sequence in output]
+
+def main():
+    while True:
+        # 사용자로부터 감정 입력 받기
+        try:
+            emotion = input("오늘의 감정을 입력하세요 (happy, sad, angry 등): ")
+            if emotion.lower() in ['happy', 'sad', 'angry']:
+                break
+            else:
+                print("입력된 감정이 잘못되었습니다. 다시 입력해주세요.")
+        except KeyboardInterrupt:
+            print("\n사용자가 입력을 취소했습니다. 프로그램을 종료합니다.")
+            return
+        
+    # 일기 생성
+    diary_entries = generate_diary(emotion)
+    # 생성된 일기 출력
+    print("오늘의 일기:")
+    for i, entry in enumerate(diary_entries, start=1):
+        print(f"{i}. {entry}")
+
+if __name__ == "__main__":
+    main()
+