luonghuyquang's picture
launch!
1336b1e
import requests
import random
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
import torch
import gradio as gr
import pandas as pd
from datetime import datetime
# Load emotion model and tokenizer
emotion_tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-emotion")
emotion_model = AutoModelForSeq2SeqLM.from_pretrained("mrm8488/t5-base-finetuned-emotion")
# Load text generation model and tokenizer
import os
token=os.getenv('hftoken')
text_tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it", token=token)
text_model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b-it", token=token)
# Set device to cpu
device = torch.device('cpu')
emotion_model.to(device)
text_model.to(device)
# Function to predict emotion
def get_emotion(text):
input_ids = emotion_tokenizer.encode(text + '</s>', return_tensors='pt').to(device)
output = emotion_model.generate(input_ids=input_ids, max_length=2)
dec = [emotion_tokenizer.decode(ids, skip_special_tokens=True) for ids in output]
label = dec[0].strip()
return label
def generate_quote(original_text, emotion):
# Generate one inspirational quote based on emotion and original text
input_text = f"Text: {original_text}\nEmotion: {emotion}\nInspirational Quote:"
input_ids = text_tokenizer(input_text, return_tensors="pt").to(device)
outputs = text_model.generate(**input_ids, max_new_tokens=70, do_sample=True, temperature=0.7)
generated_text = text_tokenizer.decode(outputs[0], skip_special_tokens=True)
if "Inspirational Quote:" in generated_text:
quote = generated_text.split("Inspirational Quote:")[1].strip().split("\n")[0]
else:
quote = generated_text.strip()
return quote
import os
import pandas as pd
# Ensure file exists and get absolute path
csv_file = os.path.join(os.getcwd(), 'diary_entries.csv')
if not os.path.exists(csv_file):
df = pd.DataFrame(columns=["Date", "Diary Text", "Emotion", "Quote"])
df.to_csv(csv_file, index=False)
else:
df = pd.read_csv(csv_file)
# Function to handle emotion detection, quote generation, and image display
def journal_interface(Diary):
try:
# Step 1: Detect Emotion
emotion = get_emotion(Diary)
# Step 2: Generate Inspirational Quote
quote = generate_quote(Diary, emotion)
# Step 3: Save to CSV
date_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
new_entry = pd.DataFrame([[date_time, Diary, emotion, quote]], columns=["Date", "Diary Text", "Emotion", "Quote"])
global df
df = pd.concat([df, new_entry], ignore_index=True)
df.to_csv(csv_file, index=False)
return emotion, quote
except Exception as e:
print(f"Error encountered: {str(e)}")
return f"Error: {str(e)}", ""
# Update the Gradio interface
interface = gr.Interface(
fn=journal_interface,
inputs=gr.Textbox(lines=5, placeholder="Enter your thoughts here..."),
outputs=[
gr.Textbox(label="Detected Emotion"),
gr.Textbox(label="Generated Quote")
],
title="AI-Powered Personal Journal",
description="Enter your thoughts, and the AI will detect the emotion and generate an inspirational quote based on it.",
theme=gr.themes.Soft()
)
# Launch the Gradio app
interface.launch(share=True)