Textual Analysis of The Office (US) Transcripts

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import seaborn as sns

import re
from nltk.corpus import stopwords
from nltk.sentiment import SentimentAnalyzer
from nltk.sentiment.vader import SentimentIntensityAnalyzer
from nltk import tokenize

from sklearn.feature_extraction.text import CountVectorizer

df = pd.read_excel("D:\\Downloads\\the-office-lines.xlsx")
df['speaker'] = [w.lower() for w in df['speaker']]

df.groupby(['speaker'])['id'].count().sort_values(ascending=False)[0:10].plot(kind='bar', figsize=(15,10))
plt.show()

def clean_words(raw_string, clean_characters=False, characters=None):
    raw_string = str(raw_string)
    cleaned = [w for w in re.sub("[^a-zA-Z]", " ", raw_string).lower().split()] 
    stops = set(stopwords.words("english"))
    cleaned = [w for w in cleaned if not w in stops]
    
    # if we want to only look at characters saying certain words (i.e. other characters)
    if clean_characters:
        cleaned = [w for w in cleaned if w in characters]
    return (" ".join(cleaned))

def gen_cleaned_script(df, clean_characters=False, characters=None):
    
    cleaned_script = []

    for i in range(0, df.shape[0]):
        if((i+1) % 5000 == 0):
            print("Review %d of %d\n" % (i+1, df.shape[0]))
        #cleaned_script.append(clean_words(df['line_text'][i]))
        cleaned_script.append(clean_words(df['line_text'][i], clean_characters, main_characters_lower))
        
    return cleaned_script

cleaned_script_all = gen_cleaned_script(df)
cleaned_script_characters = gen_cleaned_script(df, True, main_characters_lower)

def get_word_counts(text):
    
    vectorizer = CountVectorizer(analyzer="word", tokenizer=None, preprocessor=None, stop_words=None, max_features=5000)

    train = vectorizer.fit_transform(text)
    train = train.toarray()
    vocab = vectorizer.get_feature_names()

    dist = np.sum(train, axis=0)

    combined_vocab = pd.DataFrame(vocab)
    combined_vocab = combined_vocab.merge(pd.DataFrame(dist), left_index=True, right_index=True)
    combined_vocab.columns = ['Word', 'Count']
    combined_vocab.sort_values(by='Count', ascending=False)
    combined_vocab.set_index('Word', inplace=True)
    
    return combined_vocab

all_vals = get_word_counts(cleaned_script_all)
all_vals.sort_values(by='Count', ascending=False)[0:10]

michael = np.array(cleaned_script_all)[[(df.loc[df.speaker == 'michael'].index.values)]]
michael_counts = get_word_counts(michael)
michael_counts.sort_values(by='Count', ascending=False)[0:20]

vals = pd.DataFrame(index=main_characters_lower, columns=main_characters_lower)

for i in main_characters_lower:
    tmp = np.array(cleaned_script_characters)[[(df.loc[df.speaker == i].index.values)]]
    vals[i] = get_word_counts(tmp)
    
vals.fillna(0, inplace=True)

sns.clustermap(vals, annot=True, figsize=(15,10))
plt.show()

sentiment = []

for i in range(0, df.shape[0]):
    if((i+1) % 5000 == 0):
        print("Review %d of %d\n" % (i+1, df.shape[0]))
    sentiment.append(sid.polarity_scores(str(cleaned_script[i])))
    
sentiment = pd.DataFrame(sentiment)
df_sentiment = df.merge(sentiment, left_index=True, right_index=True)

main_characters = ['Michael', 'Jim', 'Dwight', 'Pam', 'Angela', 'Toby', 'Phyllis', 'Andy', 'Oscar', 'Kevin',
                  'Meredith', 'Creed', 'Kelly', 'Ryan']
main_characters_lower = [w.lower() for w in main_characters]
out = df_sentiment.loc[df_sentiment.speaker.isin(main_characters)].groupby(['season', 'speaker'])['compound'].sum().unstack()
out.fillna(0, inplace=True)
corrmat = np.corrcoef(out.T)
corrmat = pd.DataFrame(corrmat, columns=out.columns, index=out.columns)

out.plot(figsize=(15,10), cmap = cm.get_cmap('gist_ncar'))
plt.ylabel('Sum of Compound Sentiment')
plt.show()

plt.figure(figsize=(15,10))
sns.clustermap(corrmat)
plt.show()

corrmat.sum(axis=1).sort_values().plot(kind='bar')
plt.show()