view demo/text2annotation.py @ 0:90155bdd5dd6

first commit

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Apr 28 14:17:15 2018

@author: Emmanouil Theofanis Chourdakis

Takes a .txt story and annotates it based on:
    
    characters, 
    places,
    saywords,
    character_lines,
    spatial_indicators,
    
@output:
    .ann file with the same name
    .json file with the extracted character lines
    
"""

import os
import argparse
from sklearn.externals import joblib
import ner
import spacy
import re
import logging
import json
from difflib import SequenceMatcher
from neuralcoref import Coref
from rel import *

def pronoun2gender(word):
    pronoun2gender = {
        'he' : 'Male',
        'him': 'Male',
        'she': 'Female',
        'her': 'Female',
        'his': 'Male',
        'hers': 'Female',
        'himself': 'Male',
        'herself': 'Female',    
    }
    
    if word in pronoun2gender:
        return pronoun2gender[word]
    else:
        return 'neutral'


logging.basicConfig(level=logging.INFO)

# given an iterable of pairs return the key corresponding to the greatest value
def argmax(pairs):
    #https://stackoverflow.com/questions/5098580/implementing-argmax-in-python    
    return max(pairs, key=lambda x: x[1])[0]

# given an iterable of values return the index of the greatest value
def argmax_index(values):
    return argmax(enumerate(values))

# given an iterable of keys and a function f, return the key with largest f(key)
def argmax_f(keys, f):
    return max(keys, key=f)

def similar(a, b):
    """ Returns string similarity between a and b """
    # https://stackoverflow.com/questions/17388213/find-the-similarity-metric-between-two-strings
    return SequenceMatcher(None, a, b).ratio()


def get_resolved_clusters(coref):
    """ Gets a coref object (from neural coref) and
        returns the clusters as words """
    
    mentions = coref.get_mentions()
    clusters = coref.get_clusters()[0]
    result = []
    for c in clusters:
        result.append([mentions[r] for r in clusters[c]])
    return result

def cluster_word(word, clusters):
    """ Gets a word and a list of clusters of mentions
    and figures out where the word matches most based on
    string similarity """
    
    similarities = []
    for rc in clusters:
        similarity = [similar(word.lower(), c.text.lower()) for c in rc]
        similarities.append(similarity)
    max_similarities = [max(s) for s in similarities]
    if max(max_similarities) > 0.75:
        return argmax_index(max_similarities)
    else:
        return -1

def quotes2dict(text):
    new_text = text
    is_open = False
    
    quote_no = 0
    quote = []
    narrator = []
    quote_dict = {}
    
    for n, c in enumerate(text):
        if c == '"' and not is_open:
            is_open = True
            quote_dict["<nline{}>.".format(quote_no)] = ''.join(narrator)
            narrator = []
            quote_no += 1
            continue
        
        elif c == '"' and is_open:
            is_open = False
            quote_dict["<cline{}>.".format(quote_no)] = ''.join(quote)
            new_text = new_text.replace('"'+''.join(quote)+'"', "<cline{}>.".format(quote_no))
            quote = []
            quote_no += 1
            continue
        
        if is_open:
            quote.append(c)
        elif not is_open:
            narrator.append(c)
            
    return new_text, quote_dict

def figure_gender(word, clusters, character_lut):
    for c in character_lut:
        if c.lower() in [w.lower() for w in word] and character_lut[c]['gender'] in ['Male', 'Female']:
            return character_lut[c]['gender']
            
    cluster_idx = cluster_word(word, clusters)
    if cluster_idx == -1:
        return 'neutral'
    genders = [pronoun2gender(c.text) for c in clusters[cluster_idx]]
    if 'Male' in genders and 'Female' not in 'genders':
        return 'Male'
    if 'Female' in genders and 'Male' not in 'genders':
        return 'Female'
    return 'neutral'

def annotate(text, 
             ner_model, 
             rel_model,
             character_lut, 
             saywords_lut, 
             spind_lut, 
             places_lut,
             do_coreference_resolution=True):
    """
        Function which annotates entities in text
        using the model in "model",
        
        returns: A ner.Document object with tokens labelled via
                 the LUTS provided and also the NER model in "model"
    """

    # Find and store character lines in a dictionary
    logging.info('Swapping character lines for character line tags')
    processed_text, quotes = quotes2dict(text)

    # Create spacy document object  from resulting text
    # Create the nlp engine
    logging.info("Loading 'en' spacy model")
    nlp = spacy.load('en')
    
    # Loading coreference model
    coref = Coref()
    
    
    # Doing coreference resolution
    if do_coreference_resolution:
        logging.info("Doing one-shot coreference resolution (this might take some time)")
        coref.one_shot_coref(processed_text)
        resolved_clusters = get_resolved_clusters(coref)
        processed_text = coref.get_resolved_utterances()[0]
        
    # Parse to spacy document
    logging.info("Parsing document to spacy")
    doc = nlp(processed_text)    
    
    # Parse to our custom Document object
    logging.info("Parsing document to our object format for Named Entity Recognition")
    mDoc = ner.Document(doc)
    
    # Label <CLINE[0-9]+> as character line
    logging.info("Labeling character lines")
    spans = [r.span() for r in re.finditer(r'<cline[0-9]+>\.', mDoc.text)]
    for span in spans:
        mDoc.assign_label_to_tokens(span[0],span[1],'Character_Line')

    # Parse using LUTs
    
    # *- Characters
    
    # Sort by number of words so that tokens with more words override
    # tokens with less words in labelling. For example if you have 
    # `man' and `an old man' as characters, the character labelled is going to
    # be `an old man' and not the included `man'.
    logging.info("Labeling characters from LUT")
    cLUT = [c.lower() for c in sorted(character_lut, key=lambda x: len(x.split()))]
    
    # Find literals in document that match a character in cLUT 
    for c in cLUT:
        spans = [r.span() for r in re.finditer(c, mDoc.text)]
        for span in spans:
            mDoc.assign_label_to_tokens(span[0],span[1],'Character')
             
    # *- Saywords
   
    # Assign labels to saywords. here saywords contain only one token. In addition
    # we check against the saywords' lemma and not the saywords itself. 
    logging.info("Labeling saywords from LUT")
    swLUT = [nlp(sw)[0].lemma_ for sw in saywords_lut]
    for sw in swLUT:
        mDoc.assign_label_to_tokens_by_matching_lemma(sw, 'Says')
        
    # *- Places
    logging.info("Labeling places from LUT")
    plLUT = [pl.lower() for pl in sorted(places_lut, key=lambda x: len(x.split()))]
    
    # Find literals in document that match a character in cLUT 
    for pl in plLUT:
        spans = [r.span() for r in re.finditer(pl, mDoc.text)]
        for span in spans:
            mDoc.assign_label_to_tokens(span[0],span[1],'Place')
            
    # *- Spatial indicators
    logging.info("Labeling spatial indicators from LUT")
    spLUT = [sp.lower() for sp in sorted(spind_lut, key=lambda x: len(x.split()))]         
    for sp in spLUT:
        spans = [r.span() for r in re.finditer(sp, mDoc.text)]
        for span in spans:
            mDoc.assign_label_to_tokens(span[0],span[1],'Spatial_Signal')
        
    logging.info("Extracting token features")
    features, labels = mDoc.get_token_features_labels()
    
    logging.info("Predicting labels")
    new_labels = ner_model.predict(features)
    

    logging.info("Assigning labels based on the NER model")
    # If a label is not already assigned by a LUT, assign it using the model
    
    #logging.info("{} {}".format(len(mDoc.tokens), len(new_labels)))
    for m, sent in enumerate(mDoc.token_sentences):
        for n, token in enumerate(sent):
            if token.label == 'O':
                token.label = new_labels[m][n]
                
    # Assign character labels
    if do_coreference_resolution:
        logging.info('Figuring out character genders')
        character_tok_sent = mDoc.get_tokens_with_label('Character')
        for sent in character_tok_sent:
            for character in sent:
                raw_string = " ".join([c.text for c in character])
                gender = figure_gender(raw_string, resolved_clusters, character_lut)
                for tok in character:
                    if gender in ['Male', 'Female']:
                        tok.set_attribute('gender', gender)
                    
    logging.info('Predicting the correct label for all possible relations in Document')
    mDoc.predict_relations(rel_model)
    
    
    return mDoc, quotes


def doc2brat(mDoc):
    """ Returns a brat .ann file str based on mDoc """
    
    # Dictionary that maps text span -> variable (to be used when 
    # adding relations )
    span2var = {}
    
    # Variable generator for entities (T in brat format)
    tvar = ner.var_generator('T')
    
    # Variable generator for relations (E in brat format)
    rvar = ner.var_generator('E')
    
    # Variable generator for attributions (E in brat format)
    avar = ner.var_generator('A')
    
    ann_str = ""
    # Extract characters in the format 
    # T1 Character START END character string
    
    labels = ['Character', 'Says', 'Place', 'Spatial_Signal', 'Character_Line']
    
    for label in labels:
        token_sentences = mDoc.get_tokens_with_label(label)
        for tlist in token_sentences:
            if len(tlist) == 0:
                continue
            
            for tokens in tlist:
                start = tokens[0].start
                end = tokens[-1].end
                txt = mDoc.text[start:end]
                var = next(tvar)
                ann_str += "{}\t{} {} {}\t{}\n".format(var, label, start, end, txt)
                if 'gender' in tokens[0].attributes:
                    ann_str += "{}\t{} {} {}\n".format(next(avar), 'Gender', var, tokens[0].attributes['gender'])
                
                span2var[(start, end)] = var
                
    # Map relations 
    for r in mDoc.relations:
        var = next(rvar)
        trigger = r.trigger
        trigger_label = trigger[0].label[2:]
        trigger_start = trigger[0].start
        trigger_end = trigger[-1].end
        trigger_var = span2var[(trigger_start, trigger_end)]
        
        # If a trigger is Spatial_Signal then the 
        # arguments are of form Trajector and Landmark
        
        if trigger_label == 'Spatial_Signal':
            arg1_label = 'Trajector'
            arg2_label = 'Landmark'


        # If a trigger is Says then the 
        # arguments are WHO and WHAT
        
        elif trigger_label == 'Says':
            arg1_label = 'WHO'
            arg2_label = 'WHAT'
        
        # Span for the first argument
        arg1_start = r.arg1[0].start
        arg1_end = r.arg1[-1].end
        
        # Variable for the first argument
        arg1_var = span2var[(arg1_start, arg1_end)]      
        
        # Span for the second argument
        arg2_start = r.arg2[0].start
        arg2_end = r.arg2[-1].end
        
        # Variable for the second argument
        arg2_var = span2var[(arg2_start, arg2_end)]        
            
        annot_line = "{}\t{}:{} {}:{} {}:{}\n".format(var, 
                                                      trigger_label, 
                                                      trigger_var, 
                                                      arg1_label,
                                                      arg1_var,
                                                      arg2_label,
                                                      arg2_var)
        
        ann_str += annot_line
        
        
            
            
    return ann_str

if __name__=="__main__":
    argparser = argparse.ArgumentParser()
    argparser.add_argument('input_path', help='.txt file to parse')
    argparser.add_argument('ner_model_path', help='.pkl file containing NER model')
    argparser.add_argument('rel_model_path', help='.pkl file containing relational model')
    argparser.add_argument('--say-lut', help='.txt file with list of saywords')
    argparser.add_argument('--char-lut', help='.txt file with known characters')
    argparser.add_argument('--place-lut', help='.txt file with known places')
    argparser.add_argument('--spatial-indicator-lut', help='.txt file with known spatial indicators')
    argparser.add_argument('--force', help='force overwrite when there is a file to be overwritten')
    argparser.add_argument('--no-coreference-resolution', action='store_true', help='omit coreference resolution step')
    
    args = argparser.parse_args()
    
    # Load text file
    with open(args.input_path) as f:
        text = " ".join(f.read().split())
        
    output_dir = os.path.dirname(args.input_path)
    output_text_path = args.input_path[:-4] + '_processed.txt'
    output_quotes_path = args.input_path[:-4] + '_quotes.json'
    output_annotation_path = args.input_path[:-4] + '_processed.ann'
        
    # Load NER model file
    ner_model = joblib.load(args.ner_model_path)
    
    # Load REL model file
    rel_model = joblib.load(args.rel_model_path)
    
    # Load saywords
    if args.say_lut:
        saylut_path = args.say_lut
    else:
        saylut_path = 'saywords.txt'
        
    with open(saylut_path) as f:
        saylut = [s for s in f.read().split('\n') if s.strip() != '']
        
    # Load places LUT
    if args.place_lut:
        placelut_path = args.place_lut
    else:
        placelut_path = 'places.txt'
        
    with open(placelut_path) as f:
        placelut = [s for s in f.read().split('\n') if s.strip() != '']
        
    # Load spatial indicators LUT
    if args.spatial_indicator_lut:
        spatial_indicator_lut_path = args.spatial_indicator_lut
    else:
        spatial_indicator_lut_path = 'spatial_indicators.txt'
        
    with open(spatial_indicator_lut_path) as f:
        spatial_indicator_lut = [s for s in f.read().split('\n') if s.strip() != '']        
        
    # Load character LUT
    if args.char_lut:
        charlut_path = args.char_lut
    else:
        charlut_path = 'characters.txt'
        
    with open(charlut_path) as f:
        
        charlist = [s for s in f.read().split('\n') if s.strip() != ''] # One character per line
        
        character_lut = {} # Stores character attributes indexed by name
        for l in charlist:
            name, attributes = l.split(':') 
            
            gender = None
            age = None
            
            for a in attributes.split(','):
                if 'male' in a:
                    gender = a
                elif a.lower() in ['young', 'old']:
                    age = a
            
            character_lut[name] = {}
            if gender:
                character_lut[name]['gender'] = gender
            if age:
                character_lut[name]['age'] = age
        
    if args.no_coreference_resolution:
        corefres = False
    else:
        corefres = True
    mDoc, quotes = annotate(text, ner_model, rel_model, character_lut, saylut, spatial_indicator_lut, placelut, corefres)
    
    annotation_text = doc2brat(mDoc)
    
    to_save = {
            output_text_path: mDoc.text,
            output_quotes_path: json.dumps(quotes),
            output_annotation_path: annotation_text
            }
    
    
    for path in to_save:
        if not os.path.exists(path) or args.force:
            with open(path, 'w') as f:
                f.write(to_save[path])               
        else:
            overwrite = input('Path {} exists, overwrite? (y/N) '.format(path))
            if overwrite[0] in ['Y', 'y']:
                with open(path, 'w') as f:
                    f.write(to_save[path])