poem_analysis / individual_analyzes.py
esocoder's picture
.
bad1f5f
raw
history blame contribute delete
No virus
2.58 kB
import streamlit as st
import os
from transformers import pipeline
from utils import read_poems_from_directory
import matplotlib.pyplot as plt
sentiment_classifier = pipeline("sentiment-analysis")
emotion_classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", return_all_scores=True)
# Directory for poems
poems_directory = "./poems"
os.makedirs(poems_directory, exist_ok=True)
def analyze_poem(poem):
sentiment = sentiment_classifier(poem)[0]
emotion_scores = emotion_classifier(poem)[0]
emotion = max(emotion_scores, key=lambda x: x['score'])['label']
return sentiment, emotion, emotion_scores
def plot_emotion_scores(emotion_scores):
emotions = [score['label'] for score in emotion_scores]
scores = [score['score'] for score in emotion_scores]
fig, ax = plt.subplots()
ax.bar(emotions, scores)
ax.set_xlabel("Emotion")
ax.set_ylabel("Score")
ax.set_title("Emotion Scores")
plt.xticks(rotation=45)
plt.tight_layout()
return fig
def analyze_individual_page():
st.header("Analyze Individual Poems")
# Sidebar for file upload and listing files
st.sidebar.title("Upload New Poem")
uploaded_file = st.sidebar.file_uploader("Choose a text file", type="txt")
if uploaded_file is not None:
with open(os.path.join(poems_directory, uploaded_file.name), "wb") as f:
f.write(uploaded_file.getbuffer())
st.sidebar.success(f"Uploaded {uploaded_file.name}")
st.sidebar.title("Available Poems")
poem_files = [f for f in os.listdir(poems_directory) if f.endswith(".txt")]
st.sidebar.write("\n".join(poem_files))
poems = read_poems_from_directory(poems_directory)
if poems:
selected_poem = st.selectbox("Select a poem to analyze", range(len(poems)), format_func=lambda i: f"Poem {i+1}")
st.subheader(f"Analyzing Poem {selected_poem + 1}")
st.text(poems[selected_poem])
sentiment, emotion, emotion_scores = analyze_poem(poems[selected_poem])
st.write(f"Sentiment: {sentiment['label']} (score: {sentiment['score']:.2f})")
st.write(f"Dominant Emotion: {emotion}")
# Plot emotion scores
fig = plot_emotion_scores(emotion_scores)
st.pyplot(fig)
else:
st.warning("No poems found in the 'poems' directory.")
def analyze_sentiment(poems):
sentiment_labels = []
for poem in poems:
sentiment = sentiment_classifier(poem)[0]
sentiment_labels.append(sentiment['label'])
return sentiment_labels