Spaces:
Runtime error
Runtime error
File size: 7,313 Bytes
9c26d5d 3962dbd 9c26d5d 11f11a1 9c26d5d 3a2c400 11f11a1 9c26d5d 091f105 9c26d5d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
from transformers import pipeline, RobertaTokenizer, RobertaForSequenceClassification
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import make_pipeline
from datasets import load_dataset
from PIL import Image, ImageDraw, ImageFont
import textwrap
import random
from diffusers import StableDiffusionPipeline
import torch
from sklearn.metrics import classification_report, accuracy_score, f1_score
from sklearn.model_selection import train_test_split # Import train_test_split
import gradio as gr
# Load the datasets for emotion detection and quotes
emotion_dataset = load_dataset("dair-ai/emotion",trust_remote_code=True)
quotes_dataset = load_dataset("Abirate/english_quotes",trust_remote_code=True)
# Prepare the Bag-of-Words Model
vectorizer = CountVectorizer()
naive_bayes_classifier = MultinomialNB()
bow_pipeline = make_pipeline(vectorizer, naive_bayes_classifier)
texts = [example['text'] for example in emotion_dataset['train']]
labels = [example['label'] for example in emotion_dataset['train']]
train_texts, test_texts, train_labels, test_labels = train_test_split(texts, labels, test_size=0.2, random_state=42)
bow_pipeline.fit(train_texts, train_labels) # Train the Bag-of-Words model
predicted_labels = bow_pipeline.predict(test_texts)
print("Bag-of-Words Model Evaluation Metrics:")
print(classification_report(test_labels, predicted_labels))
print("Accuracy:", accuracy_score(test_labels, predicted_labels))
print("F1 Score:", f1_score(test_labels, predicted_labels, average='weighted'))
# Load the emotion classification models
distilbert_classifier = pipeline('text-classification', model='bhadresh-savani/distilbert-base-uncased-emotion')
roberta_tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
roberta_model = RobertaForSequenceClassification.from_pretrained('roberta-base')
# Assuming there's a test split for evaluation
test_data = emotion_dataset['test']
test_texts = [example['text'] for example in test_data]
test_labels = [example['label'] for example in test_data]
label_mapping = {0: 'sadness', 1: 'joy', 2: 'love', 3: 'anger', 4: 'fear', 5: 'surprise'}
test_labels = [label_mapping[label].lower() for label in test_labels]
'''
# Evaluate DistilBERT
distilbert_predictions = [distilbert_classifier(text)[0]['label'].lower() for text in test_texts]
print("DistilBERT Model Evaluation Metrics:")
print("Accuracy:", accuracy_score(test_labels, distilbert_predictions))
print("F1 Score:", f1_score(test_labels, distilbert_predictions, average='weighted'))
# Evaluate RoBERTa
roberta_predictions = []
for text in test_texts:
inputs = roberta_tokenizer(text, return_tensors="pt", padding=True, truncation=True)
outputs = roberta_model(**inputs)
prediction = torch.argmax(outputs.logits, dim=-1).item()
roberta_predictions.append(label_mapping[prediction].lower())
print("RoBERTa Model Evaluation Metrics:")
print("Accuracy:", accuracy_score(test_labels, roberta_predictions))
print("F1 Score:", f1_score(test_labels, roberta_predictions, average='weighted'))
'''
# Function to generate an image from a prompt using Stable Diffusion
def generate_image(prompt):
try:
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", variant="fp16", torch_dtype=torch.float32)#.to("cuda")
output = pipe(prompt=prompt)
image = output.images[0]
image.save("output_image.png")
return "output_image.png"
except Exception as e:
print(f"Failed during the image generation process: {e}")
return None
# Function to create an image with the quote
def create_image_with_quote(quote, author, emotion):
prompt = f"{emotion} themed image"
img_path = generate_image(prompt)
if not img_path:
print("Image generation failed.")
return None
img = Image.open(img_path)
d = ImageDraw.Draw(img)
quote_font = ImageFont.load_default()
author_font = ImageFont.load_default()
quote_wrapped = textwrap.fill(quote, width=40)
author_wrapped = textwrap.fill(f"- {author}", width=40)
draw_text_with_background(d, quote_wrapped, author_wrapped, quote_font, author_font,img.width, img.height)
return img
def draw_text_with_background(d, quote, author, quote_font, author_font, img_width, img_height):
# Calculate text bounding boxes
quote_bbox = d.textbbox((0, 0), quote, font=quote_font)
author_bbox = d.textbbox((0, 0), author, font=author_font)
# Calculate the position for the quote
quote_x = (img_width - (quote_bbox[2] - quote_bbox[0])) / 2
quote_y = (img_height - (quote_bbox[3] - quote_bbox[1])) / 2 - 20 # Slightly above the center
# Calculate the position for the author
author_x = (img_width - (author_bbox[2] - author_bbox[0])) / 2
author_y = quote_y + (quote_bbox[3] - quote_bbox[1]) + 10 # Just below the quote
# Draw background for quote
d.rectangle([quote_x - 10, quote_y - 5, quote_x + (quote_bbox[2] - quote_bbox[0]) + 10, quote_y + (quote_bbox[3] - quote_bbox[1]) + 5], fill=(255, 255, 255, 128))
# Draw background for author
d.rectangle([author_x - 10, author_y - 5, author_x + (author_bbox[2] - author_bbox[0]) + 10, author_y + (author_bbox[3] - author_bbox[1]) + 5], fill=(255, 255, 255, 128))
# Draw text over the boxes
d.text((quote_x, quote_y), quote, font=quote_font, fill="black")
d.text((author_x, author_y), author, font=author_font, fill="black")
def predict_emotion(text, model_choice):
if model_choice == 'distilbert':
result = distilbert_classifier(text)
return result[0]['label'].lower()
elif model_choice == 'bow':
prediction = bow_pipeline.predict([text])[0]
return 'positive' if prediction == 1 else 'negative'
elif model_choice == 'roberta':
inputs = roberta_tokenizer(text, return_tensors="pt", padding=True, truncation=True)
outputs = roberta_model(**inputs)
prediction = torch.argmax(outputs.logits, dim=-1)
return 'positive' if prediction.item() == 1 else 'negative'
def evaluate_quotes_for_emotion(emotion):
# Filter quotes based on the predicted emotion
suitable_quotes = [q for q in quotes_dataset['train'] if emotion.lower() in q['tags']]
if not suitable_quotes:
suitable_quotes = quotes_dataset['train'] # fallback to any quote if no tags match
selected_quote = random.choice(suitable_quotes)
return selected_quote['quote'], selected_quote.get('author', 'Unknown')
def predict_emotion_and_generate_quote(feelings, model_choice):
emotion = predict_emotion(feelings, model_choice)
quote, author = evaluate_quotes_for_emotion(emotion)
image = create_image_with_quote(quote, author, emotion)
return quote, author, image
iface = gr.Interface(fn=predict_emotion_and_generate_quote,
inputs=["text", "text"],
outputs=["text", "text", "image"],
title="Quote Generator: Feeling's Inspired",
description="Enter your feelings and choose a model to receive an inspiring quote with an accompanying image. Model Choices include distilbert,roberta,bow",
allow_flagging=False,
theme="default")
iface.launch()
|