VisionFlow / script_analyzer.py
Miguel Castro
Explicitly add [PAD] token, add Geist icons and styles
ecff187
import gradio as gr
from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer, AutoModelForCausalLM
import matplotlib.pyplot as plt
# Load sentiment and generation models as previously set up
tokenizer_sentiment = AutoTokenizer.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
model_sentiment = AutoModelForSequenceClassification.from_pretrained("cardiffnlp/twitter-roberta-base-sentiment")
sentiment_classifier = pipeline("sentiment-analysis", model=model_sentiment, tokenizer=tokenizer_sentiment)
tokenizer_gpt = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-2.7B")
model_gpt = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-2.7B")
# Updated generate_text function with padding token fix
def generate_text(prompt):
# Add '[PAD]' as a padding token if not set
if tokenizer_gpt.pad_token is None:
tokenizer_gpt.add_special_tokens({'pad_token': '[PAD]'})
# Prepare the input tensors with attention_mask and padding
inputs = tokenizer_gpt(prompt, return_tensors="pt", padding=True, truncation=True, max_length=50)
# Generate text using max_new_tokens instead of max_length
outputs = model_gpt.generate(
inputs["input_ids"],
attention_mask=inputs["attention_mask"],
max_new_tokens=50,
num_return_sequences=1,
no_repeat_ngram_size=2,
pad_token_id=tokenizer_gpt.pad_token_id # Ensure pad_token_id is set
)
generated_text = tokenizer_gpt.decode(outputs[0], skip_special_tokens=True)
return generated_text.strip()
# Function to analyze each line of the script
def analyze_script(script):
global all_scores, descriptions, music_cues
lines = script.strip().split("\n")
all_scores = []
descriptions = []
music_cues = []
analysis_results = []
for i, line in enumerate(lines):
result = sentiment_classifier(line)[0]
sentiment = result['label']
score = result['score']
description_prompt = f"Describe a scene with the sentiment '{sentiment}' for the line: '{line}'"
description = generate_text(description_prompt)
music_cue_prompt = f"Suggest music elements (like tempo, key, and instrumentation) that would fit a scene with the sentiment '{sentiment}': '{line}'"
music_cue = generate_text(music_cue_prompt)
all_scores.append(score)
descriptions.append(description)
music_cues.append(music_cue)
analysis_results.append(
{
"Line": f"Line {i + 1}: {line}",
"Sentiment": f"{sentiment} (Score: {round(score, 2)})",
"Description Suggestion": description,
"Music Cue": music_cue
}
)
graph_path = generate_script_graph()
return analysis_results, graph_path
# Generate the emotional arc graph with music cues for the entire script
def generate_script_graph():
plt.figure(figsize=(12, 6))
plt.plot(all_scores, marker='o', linestyle='-', color='b', label='Sentiment Intensity')
for i, score in enumerate(all_scores):
plt.text(i, score, music_cues[i], fontsize=8, ha='right', rotation=45)
plt.title('Emotional and Musical Arc for Entire Script')
plt.xlabel('Script Lines (Accumulative)')
plt.ylabel('Sentiment Intensity')
plt.legend()
plt.tight_layout()
plot_path = "script_emotional_arc.png"
plt.savefig(plot_path)
plt.close()
return plot_path
# Custom Gradio component to display dashboard results with icons
def format_dashboard(results):
formatted_results = ""
for result in results:
formatted_results += f"""
<div class="dashboard-box">
<p><img src="https://geist-ui.dev/icons/activity.svg" class="dashboard-icon" alt="Line icon"> <strong>{result['Line']}</strong></p>
<p><img src="https://geist-ui.dev/icons/bar-chart.svg" class="dashboard-icon" alt="Sentiment icon"> <strong>Sentiment:</strong> {result['Sentiment']}</p>
<p><img src="https://geist-ui.dev/icons/eye.svg" class="dashboard-icon" alt="Description icon"> <strong>Description Suggestion:</strong> {result['Description Suggestion']}</p>
<p><img src="https://geist-ui.dev/icons/music.svg" class="dashboard-icon" alt="Music Cue icon"> <strong>Music Cue:</strong> {result['Music Cue']}</p>
</div>
"""
return formatted_results
# Gradio interface to analyze script and display the dashboard
with gr.Blocks(css="custom.css") as interface:
gr.Markdown("## Script Sentiment and Music Cue Analyzer", elem_id="title")
gr.Markdown("Enter your script line-by-line, and this tool will analyze sentiment, generate scene descriptions, suggest music cues, and show an emotional and musical arc.", elem_id="description")
script_input = gr.Textbox(lines=10, placeholder="Enter your script here, one line per thought or dialogue.", label="Script")
display_dashboard_button = gr.Button("Analyze Script")
output_dashboard = gr.HTML(label="Dashboard Results")
output_graph = gr.Image(label="Emotional and Musical Arc for Entire Script")
def display_dashboard(script):
analysis_results, graph_path = analyze_script(script)
dashboard_content = format_dashboard(analysis_results)
return dashboard_content, graph_path
display_dashboard_button.click(display_dashboard, inputs=script_input, outputs=[output_dashboard, output_graph])
interface.launch()