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():
    # 사용자로부터 감정 입력 받기
    emotion = input("오늘의 감정을 입력하세요 (happy, sad, angry 등): ")
    # 일기 생성
    diary_entries = generate_diary(emotion)
    # 생성된 일기 출력
    print("오늘의 일기:")
    for i, entry in enumerate(diary_entries, start=1):
        print(f"{i}. {entry}")

if __name__ == "__main__":
    main()