# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] deletable=true editable=true
# # Memory Representation in Dialogue Systems
#
# The following notebook is the result of an NLP project that explores the question, "How could interaction be stored in memory, and how can that information be leveraged for further use?" 
#
# [Dialog systems](https://en.wikipedia.org/wiki/Dialog_system) can be quite useful, but have difficulty keeping track of concepts and entities dynamically. Commercial implementations among the likes of Siri, Google Assistant, and Alexa are great for performing simple tasks, but fall short when remembering ad-hoc relationships that regularly present themselves in conversation. For more information on dialogue systems, graph databases, and ontologies as they relate to this project, see the white paper entitled [IPA_Memory](/files/docs/IPA_Memory_Dan_Kondratyuk_2017.04.30.pdf) under the `docs` directory of this repository.
#
# To enhance the capabilities of dialogue systems, this notebook will provide a simple software implementation of a model that is intended to by dynamic, incremental, flexible, and interpretable. By forming high-level concepts that evolve over time, this model will evaluate the dialogue system's ability to understand user input. This notebook will show how such a system can update its internal state based on natural language facts, and retrieve results based on natural language questions. See the white paper for more details on the rationale behind these design decisions.
#
# The code below is written in Python, and uses a [Neo4j Graph Database](https://neo4j.com/product/) to provide non-volatile storage and efficient querying capabilities.
#
# The test corpus is supplied by the [bAbI Tasks Data 1-20 (v1.2)](https://research.fb.com/downloads/babi/). It contains sequences of English sentences to provide the system knowledge of a simple domain involving characters moving to different rooms and interacting with objects. Questions are inserted periodically to evaluate that the system is keeping track of these relationships accurately.
#
# ## Prerequisites to Running this Notebook
# - [Python](https://www.python.org/downloads/) (3.5+)
# - Python packages (install via pip): `pandas`, `numpy`, `nltk`, `scikit-learn`, `neo4j-driver`
# - [Neo4j](https://neo4j.com/download/) (3.1+)

# + [markdown] deletable=true editable=true
# # Part 1: bAbI QA 1
#
# ## Process the Text
#
# ### Import DataFrames
# First we will use `pandas` to import `qa1_single-supporting-fact_train.txt` from our corpus into a DataFrame called `data`. Every line in this document represents one sentence, which will be split using `nltk`'s word tokenizer.

# + deletable=true editable=true
# Import the necessary packages
import pandas as pd
import numpy as np
import nltk
from sklearn.metrics import accuracy_score

# + deletable=true editable=true
# Download NLTK packages
# An OS window should pop up for you to download the appropriate packages
# Select all-nltk and click on the download button. Once download has finished exit the window and continue.
nltk.download()

# + deletable=true editable=true
# Read the bAbI data as CSV
filename = 'resources/qa1_single-supporting-fact_train.txt'
data_qa1 = pd.read_csv(filename, delimiter='\t', names=['sentence', 'answer', 'factid'])
data_qa1 = data_qa1.fillna('')

# + [markdown] deletable=true editable=true
# The cell below shows what the input data looks like. Every `sentence` in this frame can either be a factual statement, or a question about the preceeding statements. Each statement describes four characters moving between six different rooms. The questions periodically ask the room in which a person is currently in, and the objective is to answer them all correctly, matching the corresponding `answer` column (it is blank if the sentence is a statement). The `factid` column indicates the index of the supporting facts for each answer, but we won't be needing it.
#
# Due to the nature of the model, training will not be necessary to answer each question. Therefore, the entire document will be used for test evaluation.

# + deletable=true editable=true
data_qa1[:6]

# + [markdown] deletable=true editable=true
# Next, we process this data frame by splitting the sentences and tagging each sentence by its type (statement or question). 

# + deletable=true editable=true
# Tag each sentence as a statement (S) or question (Q)
tag_sentence = lambda row: 'S' if row.answer == '' else 'Q'
data_qa1['type'] = data_qa1.apply(tag_sentence, axis=1)

# Use NLTK to tokenize the sentences into arrays of words
# If you get an error here, make sure you have downloaded the NLTK packages above
tokenize = lambda row: nltk.word_tokenize(row.sentence)[1:]
data_qa1.sentence = data_qa1.apply(tokenize, axis=1)

# Drop the factid column, as we won't need it
data_qa1 = data_qa1.drop('factid', axis=1)

# + deletable=true editable=true
data_qa1[:6]


# + [markdown] deletable=true editable=true
# We further split the `data_qa1` DataFrame into `statements` and `questions` DataFrames for easy access to all statements and questions respectively.

# + deletable=true editable=true
# Create a DataFrame with just the statements
def statements(df):
    return df[df.type == 'S'] \
        .reset_index(drop=True) \
        .drop('answer', axis=1) \
        .drop('type', axis=1)

# Create a DataFrame with just the questions
def questions(df):
    return df[df.type == 'Q'] \
        .reset_index(drop=True) \
        .drop('type', axis=1)


# + deletable=true editable=true
statements(data_qa1)[:4]

# + deletable=true editable=true
questions(data_qa1)[:2]

# + [markdown] deletable=true editable=true
# ### Extract Entities
# Next, we will extract the relevant entities from each statement and question so that we can more easily reason with these sentences.
#
# #### POS Tagging
# To process each sentence and produce a useful statement or question object, all that is necessary (for this dataset) is to use a part-of-speech tagger. The generated frame below displays the tagged list of (token, word) pairs.

# + deletable=true editable=true
# Tag each token as a part of speech
pos_tag = lambda row: nltk.pos_tag(row.sentence)
data_qa1['tag'] = data_qa1.apply(pos_tag, axis=1)

# + deletable=true editable=true
data_qa1[['sentence', 'tag']][:5]


# + [markdown] deletable=true editable=true
# #### Statements
# Due to the simplicity of the data, each statement can be thought of as a `(subject, relation, object)` triple. We would like to define a function called `extract_statement`, that when given a sequence of statement tokens, produces this triple. For instance,
# ```
# extract_statement([Mary, moved, to, the, bathroom, .]) = (Mary, moved, bathroom).
# ```
# This allows one to construct a graph of relationships between objects, as we will see in the next sections. 
#
# We can use the POS tags in the sentence to achieve this. If there is a word tagged as a proper noun, it is the subject, if there's a verb, it is the relation, and if there's a simple noun, it is the object.

# + deletable=true editable=true
def extract_statement(tags):
    '''Extracts a (subject, relation, object) triple from each statement based on the POS tags'''
    subject, relation, obj = '', '', ''
    for word,tag in tags:
        if tag == 'NNP':
            subject = word
        elif tag == 'VBD' or word == 'journeyed': # TODO: 'journeyed' is tagged improperly
            relation = word
        if tag == 'NNP' or tag == 'NN':
            obj = word
    return (subject, relation, obj)


# + [markdown] deletable=true editable=true
# #### Questions
# To test the graph, we would like to define another function `extract_question`, that when given a sequence of quesiton tokens, returns the entity that the question is asking for.
# ```
# extract_question([Where, is, Mary, ?]) = Mary
# ```
#
# The result is the subject we are querying for, whose query should return us a room to answer the question.

# + deletable=true editable=true
def extract_question(tags):
    '''Extracts the entity under discussion from each question based on the POS tags'''
    entityUnderDiscussion = ''
    # This will find the last noun in the sentence
    for word,tag in tags:
        if tag == 'NNP' or tag == 'NN':
            entityUnderDiscussion = word
    return entityUnderDiscussion


# + [markdown] deletable=true editable=true
# Then, call the appropriate function on each `DataFrame` row to extract the corresponding info.

# + deletable=true editable=true
def extract(row):
    '''Extracts the appropriate data given a processed DataFrame row'''
    if row.type == 'S':
        return extract_statement(row.tag)
    else:
        return extract_question(row.tag)


# + deletable=true editable=true
data_qa1['extracted'] = data_qa1.apply(extract, axis=1)

# + deletable=true editable=true
data_qa1[['sentence', 'extracted']][:5]


# + [markdown] deletable=true editable=true
# Voila, extraction is complete.

# + [markdown] deletable=true editable=true
# ### Debug Functions
#
# These are handy debugging functions that we will use for evaluation.

# + [markdown] deletable=true editable=true
# This function finds all statements that refer to a person.

# + deletable=true editable=true
def person_statements(person):
    '''Get all statements that refer to the specified person'''
    stat = statements(data_qa1)
    return stat[stat.extracted.map(lambda t: t[0] == person)]


# + [markdown] deletable=true editable=true
# For instance, we can find all statements that refer to Sandra.

# + deletable=true editable=true
person_statements('Sandra')[:3]


# + [markdown] deletable=true editable=true
# This function finds the `n` most recent statements that refer to a person.

# + deletable=true editable=true
def person_statements_recent(person, n=5):
    '''Get the n most recent statements that refer to the specified person in reverse chronological order'''
    return person_statements(person)[-n:].iloc[::-1]


# + [markdown] deletable=true editable=true
# For instance, we can find the 3 most recent statements Daniel has been referred in.

# + deletable=true editable=true
person_statements_recent('Daniel', n=3)

# + [markdown] deletable=true editable=true
# ## Build the Graph
# Once we have processed the data into triples, we can build graphs from them. Below we have defined a couple functions to reset the database and run queries. We will use Neo4j's Python driver to accomplish this. Note that if the URL or auth credentials of your Neo4j server are different, you will need to change them below.

# + deletable=true editable=true
from neo4j.v1 import GraphDatabase, basic_auth

# + deletable=true editable=true
# Create a neo4j session
# NOTE: Make sure that URL/credentials are correct and that Neo4j is running
driver = GraphDatabase.driver('bolt://localhost:7687', auth=basic_auth('neo4j', 'neo4j'))


# + deletable=true editable=true
# WARNING: This function will clear the database when run!
# Make sure all important data is backed up before continuing
def reset_db():
    '''Remove all nodes and relationships from the database'''
    session = driver.session()
    session.run('MATCH (n) DETACH DELETE n')


# + deletable=true editable=true
def create(query, n=0):
    '''Given a query, create a graph based on each triple in the extracted statements'''
    session = driver.session()
    stat = statements(data_qa1)
    n = len(stat) if n <= 0 else n # Run the first n statements if specified
    for subject,relation,obj in stat[:n].extracted:
        session.run(query, subject=subject, relation=relation, obj=obj)


# + [markdown] deletable=true editable=true
# ### V1: Direct relationships
# One of the first impulses when building the graph may be to represent the subject and object as nodes, and the relations as edges between them.

# + deletable=true editable=true
reset_db() # This will clear the database!

# + deletable=true editable=true
# Create a direct relationship between subject and object
v1_query = '''
    MERGE (s:SUBJECT {name: $subject}) 
    MERGE (o:OBJECT  {name: $obj}) 
    MERGE (s)-[r:RELATION {name: $relation}]->(o)
'''

create(v1_query)

# + [markdown] deletable=true editable=true
# Run the query below and see what the graph looks like. Pop open a new tab in the Neo4j browser (default http://localhost:7474/browser/) and run the query:
# ```
# MATCH (n) RETURN n LIMIT 50
# ```
# The graph is a reasonable first start, as the relations point each person to where they have been. But this poses a potential problem: how do we know where each person is right now, or where they have been previously? All we can know from the graph is which rooms a person has been in, because they may have visited them all multiple times.

# + [markdown] deletable=true editable=true
# <img src="screenshots/simple-relation.png" style="width:700px">

# + [markdown] deletable=true editable=true
# ### V2: Nodes for relationships
# One approach is to form a linked list of "events". Each event corresponds to a person updating the room that they are in. Since we chose edges to be our relations, we cannot form edges between relations. To alleviate this, we can transform the relation to a node, and draw two edges to form a 3-node triple.

# + deletable=true editable=true
reset_db() # This will clear the database!

# + deletable=true editable=true
# Represent each relation as a node
v2_query = '''
    MERGE (s:SUBJECT {name: $subject})
    MERGE (o:OBJECT  {name: $obj})
    CREATE (s)-[:R0]->(r:RELATION {name: $relation})-[:R1]->(o)
'''

create(v2_query)

# + [markdown] deletable=true editable=true
# Run the query again and see what changed. This is better, since we can see how often a room has been visited, but still doesn't solve the question as to which room a person is in at any given time.

# + [markdown] deletable=true editable=true
# ### V3: Linked list of relationships
# The final step is to build the linked list based on the order in which the relations were created. This will allow us to not only find the room a person is in right now, but produce a list of rooms that they were in, in the order that they were visited.

# + deletable=true editable=true
reset_db()

# + deletable=true editable=true
# Represent each relation as a node, ordered by a linked list (per subject)
v3_query = '''
    MERGE (s:SUBJECT {name: $subject})
    MERGE (o:OBJECT  {name: $obj})
    
    WITH s,o
    
    // Create an new relation between the subject and object
    CREATE (s)-[:R0]->(r:RELATION {name: $relation})-[:R1]->(o)
    CREATE (s)-[h:HEAD]->(r) // Make the newly created relation the head of the list
    
    WITH s,r,o,h
    
    // Find the previous head of the list (if none exist, this query will terminate here)
    MATCH (s)-[h_prev:HEAD]->(r_prev:RELATION)
    WHERE h_prev <> h
    
    // Complete the link, remove the previous head pointer
    CREATE (r_prev)-[:NEXT]->(r)
    DELETE h_prev
'''

# + deletable=true editable=true
session = driver.session()
# Create an index for faster access
session.run('CREATE INDEX ON :SUBJECT(name)')
session.run('CREATE INDEX ON :RELATION(name)')
session.run('CREATE INDEX ON :OBJECT(name)')
create(v3_query)


# + [markdown] deletable=true editable=true
# Check the new graph out and see what changed. It's helpful to change the colors of the nodes and edges to visualize this better.

# + [markdown] deletable=true editable=true
# <img src="screenshots/local-list.png" style="width:800px">

# + [markdown] deletable=true editable=true
# ## Query the Graph
# Now we can ask the graph useful questions.
# #### Find the room a person is in

# + deletable=true editable=true
def find_person(person):
    '''Find the room a person is currently in'''
    query = '''
        MATCH (s:SUBJECT {name:$name})-[:HEAD]->(r:RELATION)-->(o:OBJECT)
        RETURN s AS subject, r AS relation, o AS obj
    '''
    return session.run(query, name=person)


# + [markdown] deletable=true editable=true
# Using the graph-querying function above we can ask, "Where is Mary?"

# + deletable=true editable=true
# Note: If this is run less than a second after creating the knowledge graph, 
# the Python driver may cause a race condition where the graph 
# isn't finished updating, which could give you the wrong answer.
session = driver.session()
record = find_person('Mary').single()
print(record['obj'].get('name'))

# + [markdown] deletable=true editable=true
# According to the graph, Mary is in the kitchen. We can verify that this is true with the debug function below, and we can see the corresponding sentence that generated the relationship as well.

# + deletable=true editable=true
person_statements_recent('Mary', n=1)


# + [markdown] deletable=true editable=true
# #### Find the rooms a person has been in (reverse chronological order)

# + deletable=true editable=true
def find_person_history(person, n=100):
    '''Find the list of rooms a person was in, ordered by recency'''
    length = str(n) if n >= 1 else ''
    
    query = '''
        MATCH (s:SUBJECT {name:$name})-[:HEAD]->(r:RELATION)-->(o:OBJECT)
        MATCH (s)-->(r_prev:RELATION)-[k*1..%s]->(r), (r_prev)-->(o_prev:OBJECT)
        
        WITH size(k) AS dist, r, o, r_prev, o_prev
        ORDER BY size(k)
        
        WITH r, o, r_prev, o_prev
        RETURN [r.name] + collect(r_prev.name) AS relation, [o.name] + collect(o_prev.name) AS obj
    '''
    query = query % length
    
    session = driver.session()
    record = session.run(query, name=person).single()
    history = list(zip(record['relation'], record['obj']))[:-1]
    
    return history


# + [markdown] deletable=true editable=true
# A more advanced question that we get for free based on the graph structure is, "Where has John been recently?"

# + deletable=true editable=true
find_person_history('John', n=5)

# + [markdown] deletable=true editable=true
# Verify that John has been to to those places, in that order.

# + deletable=true editable=true
person_statements_recent('John', n=5)


# + [markdown] deletable=true editable=true
# #### Find the history of visitors for a room

# + deletable=true editable=true
def find_room_visitors(room):
    '''Find the list of visitors a room has, ordered by recency'''
    
    query = '''
        MATCH (r:RELATION)-->(o:OBJECT {name:$name})
        RETURN count(r) AS count
    '''
    
    session = driver.session()
    record = session.run(query, name=room).single()
    
    return record['count']


# + [markdown] deletable=true editable=true
# Just for fun, we can find out how many times a room has been visited. "How many times has the office been visited?"

# + deletable=true editable=true
find_room_visitors('office')


# + [markdown] deletable=true editable=true
# ## Calculate an Accuracy Score

# + deletable=true editable=true
def get_answers(row):
    '''Given an input row merge the statement in the graph, 
    or query the graph if it is a question'''
    if row.type == 'S':
        subject,relation,obj = row.extracted
        session.run(v3_query, subject=subject, relation=relation, obj=obj)
        return ''
    elif row.type == 'Q':
        person = row.extracted
        # WARNING: do not consume the result (e.g., call .consume() or .single()) 
        # until the entire iteration is done.
        # Failure to do so may cause the queries to be VERY slow!
        return find_person(person)


# + [markdown] deletable=true editable=true
# Start all over, and run through the entire dataset.

# + deletable=true editable=true
reset_db()

# + deletable=true editable=true
session = driver.session()
results = data_qa1.apply(get_answers, axis=1)
results = [x for x in results if x != '']
predicted = [result.single()['obj'].get('name') for result in results]

# + [markdown] deletable=true editable=true
# The `predicted` array contains the predicted answer to each question.`

# + deletable=true editable=true
predicted[:5]

# + [markdown] deletable=true editable=true
# The `actual` array contains the actual answers to all questions.

# + deletable=true editable=true
actual = list(data_qa1[data_qa1.type == 'Q'].answer)

# + deletable=true editable=true
actual[:5]

# + deletable=true editable=true
accuracy_score(actual, predicted)

# + [markdown] deletable=true editable=true
# And just like that, we get an accuracy of 100%. Of course, this dataset is very simple (and machine generated), so it should be of no surprise. But one notable achievement is that the graph we created can generalize to any statements of the form, `(subject, relation, object)`.

StarcoderdataJupyterScriptsDedupFiltered