extract_relationship_script.py

# notebook procedures wrapped in easier to use functions

import spacy 
from spacy import displacy
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import networkx as nx
from tqdm import tqdm
import os
from pyvis.network import Network

# Function to import the books as a list of files
def import_books(folder_path: str):
    """ Function to import the books as a list of files
    folder_path: The path were books to process are.
    Output: a list of files """
    files = [f.path for f in os.scandir(folder_path) if f.name.endswith('.txt')]
    return files

# Function to filter the entities found in a sentence
def filter_entity(entity_list, characters_df):
    """
    Function to filter the entities found in a sentence
    Input: entity_list = the list to filter, characters_df = elements to be kept
    Output: filtered entity_list
    """
    filtered_list = []
    for entity in entity_list:
        if entity in characters_df['Name'].values:
            filtered_list.append(entity)
        elif entity in characters_df['Alias'].values:
            # If the entity is an alias, i use the name instead
            corresponding_name = characters_df[characters_df['Alias'] == entity]['Name'].values[0]
            filtered_list.append(corresponding_name)
    return filtered_list

def character_dataframe(df_path: str):
    """
    Import as pandas dataframe the csv file and keep only the first name  
    Input: df_path = the path of the csv file
    Output: pandas dataframe"""
    # Read the characters dataframe generated by the "character_scraping.ipynb" script
    characters_df = pd.read_csv(df_path)
    # Knowing that in many cases, in the book, the characters are referred by their alias or only the first name,
    # In the name column i consider only the first name of the character
    characters_df['Name'] = characters_df['Name'].apply(lambda x: x.split(' ')[0])
    return characters_df

# given a book and a dataframe of characters, extract the relationships from the book returning an edgelist
def extract_relationships(book, characters_df, w_size=5):
    """
    Given a book and a dataframe of characters, extract the relationships from the book returning an edgelist
    Input: book = the book file, characters_df = dataframe of characters, w_size = (default = 5) the size of the moving window
    Output: edgelist (pd.DataFrame)
    """
    window_size = w_size
    relationships = []

    file_text = open(book, encoding='utf-8').read()
    print('book opened')
    NER = spacy.load("en_core_web_sm")
    NER.max_length = len(file_text)
    # Give the text in input to the NER model, output will be {text: entities}
    file_doc = NER(file_text[:50000])
    print('book inspected')

    # Extract the entities from the text

    # List of dictionaries {sentence: entities in the sentence}
    sent_entity_df = []
    for sent in file_doc.sents:
        entity_list = [ent.text for ent in sent.ents]
        sent_entity_df.append({"sentence": sent, "entities": entity_list})
    
    sent_entity_df = pd.DataFrame(sent_entity_df)
    print(sent_entity_df.head(3))
    # Apply the filter function to the entities found in each sentence
    sent_entity_df['characters_entities'] = sent_entity_df['entities'].apply(lambda x: filter_entity(x, characters_df))
    print(sent_entity_df.head(3))
    # Filter the sentences which don't contain any character
    sent_entity_filtered = sent_entity_df[sent_entity_df['characters_entities'].map(len) > 0]
    print(sent_entity_filtered)
    window_size = w_size
    relationships = []

    # iterate over all the sentences
    for i in range(0, sent_entity_filtered.index[-1]):
        # to not have an index out of range errors at the end
        end_i = min(sent_entity_filtered.index[-1],i + window_size)
        # put all the characters that appear in the window in a list (sum of lists is concatenation)
        char_list = sum((sent_entity_filtered['characters_entities'][i:end_i]), [])
        #remove relationship between the same character (cast to set to remove duplicates and then back to list)
        char_unique = list(set(char_list))
        # if there are more than 1 character in the window, create all the sequential relationships
        if len(char_unique) > 1:
            for index, char_a in enumerate(char_unique[:-1]):
                char_b = char_unique[index + 1]
                relationships.append({'source': char_a, 'target': char_b})

    # Create dataframe using the list of relationships
    relationships_df = pd.DataFrame(relationships)

    # sort the relationships b->a and a->b to have a unique relationship for each pair of characters
    relationships_df = pd.DataFrame(np.sort(relationships_df, axis=1), columns=relationships_df.columns)
    relationships_df['value'] = 1

    relationships_df = relationships_df.groupby(['source', 'target'], sort=False, as_index=False).sum()
    # a dataframe witch contains source, target and value can be called a edge list
    # we want that alias and name are considered as the same character

    # drop the rows where value is less than 5
    relationships_df = relationships_df[relationships_df['value'] >= 5]

    return relationships_df

# given an edgelist, create a graph
def create_graph(relationships_df):
    """
    Input: relationships_df = edgelist for the graph
    Output: NetworkX graph
    """
    Graph = nx.from_pandas_edgelist(relationships_df, 'source', 'target', edge_attr= 'value', create_using=nx.Graph())
    return Graph

# given a graph, extract the degree centrality of every character
def extract_degree_from_graph(Graph):
    """
    Given a graph, extract the degree centrality of every character
    Input: Graph = NetworkX graph
    Output: pd.DataFrame containing Name:degree_centrality
    """
    degree_centrality = nx.degree_centrality(Graph)
    degree_centrality_df = pd.DataFrame.from_dict(degree_centrality, orient='index', columns=['degree'])
    degree_centrality_df.reset_index().rename(columns={'index': 'Name'})
    return degree_centrality_df

# given a graph, extract the betweeness centrality of every character
def extract_betweeness_from_graph(Graph):
    """
    Given a graph, extract the betweeness centrality of every character
    Input: Graph = NetworkX graph
    Output: pd.DataFrame containing Name:betweeness_centrality
    """
    betweeness_centrality = nx.closeness_centrality(Graph)
    betweeness_centrality_df = pd.DataFrame.from_dict(betweeness_centrality, orient='index', columns=['betweeness'])
    betweeness_centrality_df.reset_index().rename(columns={'index': 'Name'})
    return betweeness_centrality_df

# given a graph, extract the closeness centrality of every character
def extract_closeness_from_graph(Graph):
    """
    Given a graph, extract the closeness centrality of every character
    Input: Graph = NetworkX graph
    Output: pd.DataFrame containing Name:closeness_centrality
    """
    closeness_centrality = nx.closeness_centrality(Graph)
    closeness_centrality_df = pd.DataFrame.from_dict(closeness_centrality, orient='index', columns=['closeness'])
    closeness_centrality_df.reset_index().rename(columns={'index': 'Name'})
    return closeness_centrality_df

def complete_df(degree_centrality_df, betweeness_centrality_df, closeness_centrality_df , characters_centralities):
    """ Merge degree, closeness, degree centralities in a pd.DataFrame """
    characters_centralities['degree'] = degree_centrality_df['degree']
    characters_centralities['betweeness'] = betweeness_centrality_df['betweeness']
    characters_centralities['closeness'] = closeness_centrality_df['closeness']
    return characters_centralities


""" 
books = import_books('D:\VS codici\Codici Python\Project GOT\Archive')
print(books)
characters = character_dataframe('D:\VS codici\Codici Python\Project GOT\character_list.csv')
print(characters.head(3))

centralities_tens = []
books_graph = []

for book in books:
    rels = extract_relationships(book, characters)
    graph = create_graph(rels)
    books_graph.append(graph)
    characters_centralities = pd.DataFrame(index = characters['Name'].values, columns = {
                                 'degree': np.zeros(len(characters['Name'])),
                                 'betweeness': np.zeros(len(characters['Name'])),
                                 'closeness': np.zeros(len(characters['Name']))
                                 }
                                 )
    
    centralities_tens.append(complete_df(extract_degree_from_graph(graph),
                                         extract_betweeness_from_graph(graph),
                                         extract_closeness_from_graph(graph),
                                         characters_centralities)) """