build-corpus

#!/usr/bin/env python3

# TechKnAcq: Build Corpus
# Jonathan Gordon

import sys
import os
import io
import string
import click
import re
import json
import requests
import nltk
import time
import random
import subprocess
import multiprocessing as mp
import wikipedia as wiki

from math import log
from functools import partial
from collections import defaultdict, Counter
from noaho import NoAho
from nltk.util import ngrams

from websearch import WebSearch

from techknacq.corpus import Corpus, Document
from techknacq.lx import ScrabbleLexicon, StopLexicon, SentTokenizer


# Parameters

PROCESSES = int(.5 * mp.cpu_count())

def read_key(fname):
    full_path = os.path.expanduser(fname)
    if os.path.exists(full_path):
        return open(full_path).read().strip()
    print('Could not read API key from', fname, file=sys.stderr)
    print('Some corpus expansion methods may fail.', file=sys.stderr)
    return None


SD_API_KEY = read_key('~/.t/sd.txt')
BING_API_KEY = read_key('~/.t/bing.txt')
GOOGLE_API_KEY = read_key('~/.t/google.txt')

###


def get_terms(corpus, n=4000):
    """Return an ordered list of technical terms or names of research
    topics, based on citation graph density."""

    stop = StopLexicon()
    scrabble = ScrabbleLexicon()

    ngrams = defaultdict(set)
    citations = defaultdict(set)

    for doc in corpus:
        for ng in get_ngrams([doc.title]):
            if good_ngram(ng, scrabble, stop):
                ngrams[ng].add(doc.id)
        for ref in doc.references:
            citations[doc.id].add(ref)
            citations[ref].add(doc.id)

    ngrams = filter_plurals(ngrams)

    ngram_counts = {x: len(ngrams[x]) for x in ngrams}
    filtered = filter_subsumed(ngram_counts)

    if citations:
        ngrams = score_ngrams(ngrams, citations)
        ngrams = filter_subsumed(ngrams)
        return [' '.join(x) for x in sorted(ngrams, key=lambda x: ngrams[x],
                                            reverse=True) if x in filtered][:n]
    else:
        return [' '.join(x) for x in sorted(filtered, key=lambda x: filtered[x],
                                            reverse=True)][:n]


def get_ngrams(l, min=1):
    """Return the set of all n-grams that occur in the input strings, at
    least 'min' times."""

    def everygrams(seq):
        """Return all possible n-grams generated from a sequence of items,
        as an iterator."""
        for n in range(1, len(seq) + 1):
            for ng in ngrams(seq, n):
                yield ng

    counts = defaultdict(int)
    for s in l:
        tokens = nltk.tokenize.word_tokenize(s)

        # Preserve case for acronyms inside n-grams but lowercase everything
        # else.
        new_tokens = []
        for token in tokens:
            if token.isupper() and len(token) > 1:
                new_tokens.append(token)
            else:
                new_tokens.append(token.lower())
        tokens = new_tokens

        for ngram in everygrams(tokens):
            counts[ngram] += 1

    return {x: counts[x] for x in counts if counts[x] >= min}


def good_ngram(ng, scrabble, stop):
    """Check if the n-gram is good: It's not a single word that would be
    found in a Scrabble dictionary, and it doesn't begin or end with a
    stopword."""

    # Remove single-word n-grams that are in the Scrabble dictionary.
    if len(ng) == 1 and ng[0].lower() in scrabble:
        if ng[0] not in ['POS', 'HMM', 'EM', 'PENMAN', 'CHILDES', 'AI']:
            return False

    if len(ng) == 1 and not ng[0][0].isalpha():
        return False

    # Remove n-grams that begin or end with stopwords, e.g., conjunctions,
    # prepositions, or personal pronouns.
    for word in [ng[0], ng[-1]]:
        if word in stop or word.lower() in stop:
            return False
        if word[-1].isdigit() and '-' in word:
            # E.g., 'COLING-92'.
            return False
        if word[-1] == '-' or word[-1] == '.':
            return False

    # Remove n-grams containing words with no ASCII letters, e.g., numbers or
    # symbols.
    for word in ng:
        if not any(c in string.ascii_letters for c in word):
            return False
        if '*' in word:
            return False
    return True


def filter_plurals(ngrams):
    """Remove regular plurals if the list includes the singular."""

    remove = set()
    for ng in ngrams:
        pl = ng[:-1] + (ng[-1] + 's',)
        if pl in ngrams:
            remove.add(pl)
            try:
                # Counts
                ngrams[ng] += ngrams[pl]
            except TypeError:
                # Sets of IDs
                ngrams[ng] |= ngrams[pl]
        pl = ng[:-1] + (ng[-1] + 'es',)
        if pl in ngrams:
            remove.add(pl)
            try:
                # Counts
                ngrams[ng] += ngrams[pl]
            except TypeError:
                # Sets of IDs
                ngrams[ng] |= ngrams[pl]
    for ng in remove:
        del ngrams[ng]
    return ngrams


def filter_subsumed(ngrams):
    """Remove n-grams whose scores are within 25% of subsuming n+1-grams."""

    remove = set()
    for ng in ngrams:
        if len(ng) == 1:
            continue
        shorter = ng[:-1]
        if shorter in ngrams and ngrams[shorter] <= ngrams[ng]*1.3:
            remove.add(shorter)
        shorter = ng[1:]
        if shorter in ngrams and ngrams[shorter] <= ngrams[ng]*1.3:
            remove.add(shorter)
    for ng in remove:
        del ngrams[ng]
    return ngrams


def score_ngrams(ngram_docs, citations, pc=0.9):
    """Score n-grams based on the density of the citation graph for
    documents containing them."""

    scores = dict()
    for ng in ngram_docs:
        # Build term citation graph.
        gt_nodes = ngram_docs[ng]
        gt_edges = {}
        for node in gt_nodes:
            dests = citations[node].intersection(gt_nodes)
            if dests:
                gt_edges[node] = dests

        # Terms and sets for the equations
        n = float(len(citations))
        na = float(len(gt_nodes))
        nca = float(len(list(gt_edges.keys())))
        connected_nodes = list(gt_edges.keys())
        other_nodes = gt_nodes - set(gt_edges.keys())

        # Skip ones where there aren't enough nodes to count.
        if na < 4.0:
            continue

        # Likelihood of observing a tightly connected term citation graph
        # if the term is a topic.
        oh1 = nca * log(pc) + (na - nca) * log(1 - pc)

        # Likelihood of observing a tightly connected term citation graph
        # if the term is not a topic.
        term1 = 0.0
        for i in connected_nodes:
            li = len(citations[i])
            term1 += log(1.0 - (1.0 - (na - 1.0)/(n - 1.0))**li)
        term2 = 0.0
        for i in other_nodes:
            li = len(citations[i])
            term2 += li * log(1.0 - (na - 1.0)/(n - 1.0))
        oh0 = term1 + term2

        scores[ng] = oh1 - oh0
    return scores


###


# ScienceDirect API documentation:
#   http://api.elsevier.com/documentation/SCIDIRSearchAPI.wadl


def search_sd(terms):
    """Search ScienceDirect for documents containing specified terms."""

    print('-- Search ScienceDirect:', len(terms), 'terms.')

    books = defaultdict(set)
    # ScienceDirect can return errors when too many requests are sent in
    # parallel.
    # for result in pool.imap(search_sd_helper, chunks(terms, 50)):
    for result in map(search_sd_helper, chunks(terms, 15)):
        for book in result:
            books[book] |= result[book]

    return books


def search_sd_helper(terms):
    books = defaultdict(set)

    query = ' OR '.join(['"'+x+'"' for x in terms])

    # Filter chaff in query.
    stop = ['glossary', 'index', 'removed', 'dedication', 'bibliography',
            'copyright', 'front matter', 'frontmatter', 'contributors',
            'acknowledgment', 'acknowledgement', 'preface', 'references',
            'postface', 'edited by', 'list of', 'about the', 'afterword',
            'abbreviations', 'further reading', 'guide to readers',
            'foreword', 'editors', 'author', 'abstract',
            'how to use this book']
    query += ' AND NOT title(' + ' OR '.join(['"'+s+'"' for s in stop]) + ')'

    url = 'http://api.elsevier.com/content/search/index:scidir'
    vals = {'query': query,
            'subscribed': True,
            'oa': True,
            'content': 'nonserial',
            'count': 100,
            'apikey': SD_API_KEY}
    r = requests.get(url, params=vals)

    try:
        entries = r.json()['search-results']['entry']
    except:
        sys.stderr.write('Got bad response to search.\n')
        print('Search parameters:')
        print(vals)
        print('Reply:', r.status_code)
        print(r.text)
        return books

    for entry in entries:
        try:
            pii = re.sub('[-()./]', '', entry['pii'])
            doc_title = entry['dc:title']
            isbn = entry['prism:isbn']
            pub_title = entry['prism:publicationName']
        except KeyError:
            continue

        books[(isbn, pub_title)].add((pii, doc_title))

    return books


def download_sd(docs_by_book):
    """Download documents from ScienceDirect."""

    doc_ids = set()
    for book in docs_by_book:
        for pii, _ in docs_by_book[book]:
            doc_ids.add(pii)

    print('-- Download ScienceDirect:', len(doc_ids), 'documents.')

    pool.map(download_sd_doc, doc_ids)

    curried = partial(download_sd_doc, view='ref')
    pool.map(curried, doc_ids)


def download_sd_doc(pii, view='full'):
    """Download the ScienceDirect document with the specified PII unless
    it has previously been downloaded."""

    file_path = os.path.join(outputdir, 'sd-download', pii+'-'+view+'.xml')
    if not os.path.exists(file_path):
        print('   Download:', pii + '-' + view + '.xml')

        url = 'http://api.elsevier.com/content/article/pii:' + pii
        vals = {'view': view,
                'apikey': SD_API_KEY}
        r = requests.get(url, params=vals)

        if r.status_code != requests.codes.ok:
            print('!! ScienceDirect server error:', r.status_code,
                  file=sys.stderr)
            print(r.text, file=sys.stderr)
            return

        with io.open(file_path, 'w', encoding='utf8') as out:
            out.write(r.text)


def filter_sd(docs_by_book, terms):
    """Find documents that are highly relevant to the list of terms based
    on the number of occurrences of terms and the number of potentially
    matching documents in the entire book."""

    num_docs = sum([len(docs_by_book[book]) for book in docs_by_book])
    print('-- Filter ScienceDirect:', num_docs, 'documents.')

    books = Counter()
    docs_high = set()
    docs_low = set()

    for book in docs_by_book:
        curried = partial(count_terms_in_doc, terms=terms)
        doc_ids = [pii for (pii, doc_title) in docs_by_book[book]]
        for pii, count, unique in pool.imap(curried, doc_ids):
            if count >= 15 and unique >= 5:
                docs_high.add(pii)
                books[book] += 1
            if count >= 4 and unique >= 2:
                docs_low.add(pii)

    print('-- Potential ScienceDirect:', len(docs_low), 'documents.')

    filtered = set()
    for (book, _) in books.most_common():
        _, book_title = book
        if len(filtered) > .75 * len(docs_high):
            break

        print()
        print(' -', book_title)

        for pii, doc_title in docs_by_book[book]:
            if pii in docs_low:
                print('  ', doc_title, pii)
                filtered.add(pii)

    print('-- ScienceDirect:', len(filtered), 'documents selected.')

    return filtered


def count_terms_in_doc(pii, terms):
    """Given a ScienceDirect PII and a list of terms, count how many times
    those terms occur in the corresponding document, total and unique."""

    # We can't pass the trie as an argument when using multiprocessing.
    trie = NoAho()
    for term in terms:
        trie.add(term)

    file_path = os.path.join(outputdir, 'sd-download', pii + '-full.xml')
    try:
        text = io.open(file_path, 'r', encoding='utf8').read().lower()
        matches = [text[x[0]:x[1]] for x in trie.findall_long(text)]
    except:
        return [pii, 0, 0]

    return [pii, len(matches), len(set(matches))]


def get_sd_docs(ids):
    sd_corpus = Corpus()

    for pii in ids:
        f = os.path.join(outputdir, 'sd-download', pii + '-full.xml')
        fxml = os.path.join(outputdir, 'sd-download', pii + '-ref.xml')
        d = Document()
        d.read_sd(f, fxml)
        sd_corpus.add(d)

    return sd_corpus


###


def search_wiki(terms):
    """Search Wikipedia for articles containing specified terms."""

    print('-- Search Wikipedia:', len(terms), 'terms.')

    titles = set()
    for result in pool.imap(search_wiki_helper, terms):
        titles.update(result)

    print('-- Download Wikipedia:', len(titles), 'articles.')

    categories = defaultdict(set)
    docs_high = []
    docs_low = []

    curried = partial(get_wiki_article, terms=terms)
    for article, count, unique in pool.imap(curried, titles):
        if article:
            if count >= 10 and unique >= 5:
                docs_high.append(article)
                try:
                    for c in article.categories:
                        categories[c].add(article.title)
                except:
                    pass
            if count >= 4 and unique >= 2:
                docs_low.append(article)
            else:
                print('   - Not enough terms:', article.title)
    print('-- Choosing categories.')

    covered = set()
    top_categories = set()
    for cat in sorted(categories, key=lambda x: len(categories[x]),
                      reverse=True):
        if len(covered) >= .75 * len(docs_high):
            break
        if len(categories[cat]) == 1:
            break
        if cat.startswith('All ') or cat.startswith('Pages ') or \
           'Wikipedia' in cat or cat.endswith(' stubs') or \
           'articles' in cat.lower() or 'category' in cat.lower() or \
           'CS1' in cat or 'US-centric' in cat or 'USA-centric' in cat or \
           'iving people' in cat or 'Wikidata' in cat or \
           cat.startswith('Use ') or cat.endswith(' alumni') or \
           'University' in cat or cat.endswith(' deaths') or \
           cat.endswith(' births') or 'ompanies' in cat or \
           'EngvarB' in cat or 'ambiguous time' in cat or \
           'needing confirmation' in cat:
            continue
        if len(top_categories) > .75 * len(categories):
            break
        print('  - Selected category:', cat,
              len(categories[cat]))
        top_categories.add(cat)
        covered |= categories[cat]

    print('-- Choosing articles.')

    good_articles = []
    for article in docs_low:
        try:
            categories = article.categories
            text = article.content.lower()
        except:
            continue

        if not top_categories.intersection(set(categories)):
            print('   - Bad categories:', article.title)
            continue
        print('   + Good:', article.title)
        good_articles.append(article)

    return good_articles


def search_wiki_helper(term):
    return set(wiki.search(term, results=2))


def get_wiki_article(title, terms):
    # We can't pass the trie as an argument when using multiprocessing.
    trie = NoAho()
    for term in terms:
        trie.add(term)

    for bad in ['Category:', 'List of ', 'Index of ']:
        if bad in title:
            return (None, None, None)

    try:
        article = wiki.page(title)
        text = article.content.lower()
    except:
        return (None, None, None)

    if len(text) < 200:
        return (None, None, None)

    matches = [text[x[0]:x[1]] for x in trie.findall_long(text)]

    return (article, len(matches), len(set(matches)))


def get_wiki_docs(articles):
    wiki_corpus = Corpus()
    for doc in pool.imap(export_wiki_article, articles):
        if doc:
            wiki_corpus.add(doc)
    return wiki_corpus


def export_wiki_article(article):
    d = Document()
    try:
        d.id = 'wiki-' + re.sub('[:/ ]', '_', article.title.lower())
        d.authors = ['Wikipedia']
        d.title = article.title
        d.book = 'Wikipedia'
        d.url = article.url
        d.references = ['wiki-' + re.sub('[:/ ]', '_', x.lower())
                        for x in article.links]
    except:
        print('Error processing article.')
        return None

    st = SentTokenizer()
    for line in article.content.split('\n'):
        if line.startswith('==') and line.endswith('=='):
            line = line.replace('Edit =', ' =')
            try:
                heading = re.match('=+ (.+) =+$', line).group(1).strip()
                d.sections.append({'heading': heading,
                                   'text': []})
            except: pass
        elif len(d.sections) == 0:
            d.sections.append({'text': st.tokenize(line)})
        else:
            d.sections[-1]['text'] += st.tokenize(line)
    return d


###


def search_web(terms):
    print('-- Search Web for technical terms.')

    tutorials = Counter()
    for results in pool.imap(search_web_term, chunks(terms, 10)):
        for tutorial in results:
            tutorials[tutorial] += 1
    for results in pool.imap(search_web_term, [[x] for x in terms[:200]]):
        for tutorial in results:
            tutorials[tutorial] += 1

    # Remove singletons.
    tutorials -= Counter({x: 1 for x in list(tutorials)})

    for t in tutorials:
        print(' -', t[0])
        print('  ', t[1][:70])

    print('--', len(tutorials), 'matches.')

    return tutorials


def search_web_term(terms):
    # Slow down to avoid overloading the service.
    time.sleep(.2)

    w = WebSearch(key=GOOGLE_API_KEY, cx='014701580703549033938:zm7x9pkz36w')

    titles = {}
    tutorials = set()

    query = 'filetype:pdf "this tutorial" "' + \
      '" OR "'.join(terms) + '"'

    exclude = [
        'advisor', 'appendix', 'assignment', 'author', 'biographical',
        'chapter', 'college', 'course', 'curriculum', 'cv', 'dissertation',
        'doctor of philosophy', 'dr', 'enrollment', 'intelligent tutoring',
        'lecture', 'ph. d', 'ph.d', 'phd', 'powerpoint', 'problem set',
        'programme', 'references', 'research statement', 'resume',
        'schedule', 'session', 'slide', 'syllabus', 'table of contents',
        'tutorial dialog', 'tutorial section', 'tutorial system',
        'university', 'vitae', 'vitæ', 'vita', 'proceedings', 'transcript',
        'conference', 'agenda', 'workshop', 'hello', 'abstracts',
        'instructor', 'manual'
    ]

    for result in w.search(query, limit=5):
        if not result or not result.url or not result.title:
            continue

        # Does it look like an ACL paper?
        if re.search(r'/[a-z][0-9][0-9]-[0-9][0-9][0-9][0-9]\.pdf',
                     result.url.lower()):
            continue

        bad = False
        for term in exclude:
            if term in result.title.lower():
                bad = True
                break
            if term in result.url.lower():
                bad = True
                break
            if term in result.description.lower():
                bad = True
                break
            if re.search('\bprogram\b', result.title.lower()):
                bad = True
                break
        if bad:
            continue

        tutorials.add(result.url)
        titles[result.url] = result.title

    return set([(titles[x], x) for x in tutorials])


def download_web(tutorials):
    print('-- Downloading', len(tutorials), 'tutorials.')

    pool.map(download_tutorial, tutorials)


def download_tutorial(tutorial):
    title, url = tutorial

    fname = 'web-' + re.sub('[ :/()]', '_', title[:40].lower())
    fname = fname.replace('_...', '')
    fname = re.sub('_+', '_', fname)
    fname = ''.join(filter(lambda x: x in string.printable, fname))
    file_path = os.path.join(outputdir, 'web-download', fname + '.pdf')

    if not os.path.exists(file_path):
        try:
            r = requests.get(url)
            with open(file_path, 'wb') as out:
                out.write(r.content)
        except:
            print(' ! Download failed.', file=sys.stderr)
            if os.path.exists(file_path):
                os.remove(file_path)


def get_web_docs(tutorials):
    web_corpus = Corpus()
    for doc in pool.imap(export_tutorial, tutorials):
        if doc:
            web_corpus.add(doc)
    return web_corpus


def export_tutorial(tutorial):
    def count_pages(f):
        r = re.compile(b"/Type\s*/Page([^s/]|$)", re.MULTILINE|re.DOTALL)
        return len(r.findall(open(f, "rb").read()))

    title, url = tutorial

    fname = 'web-' + re.sub('[ :/()]', '_', title[:40].lower())
    fname = fname.replace('_...', '')
    fname = re.sub('_+', '_', fname)
    fname = ''.join(filter(lambda x: x in string.printable, fname))

    print('- Export:', fname + '.json')

    pdfpath = os.path.join(outputdir, 'web-download', fname + '.pdf')

    if not os.path.exists(pdfpath):
        print(' ! Missing download:', fname, file=sys.stderr)
        return

    pages = count_pages(pdfpath)
    if pages in [1, 2] or pages > 100:
        print(' ! Inappropriate number of pages:', fname, file=sys.stderr)
        return

    try:
        text = subprocess.check_output(['pdftotext', pdfpath, '-'], timeout=10)
    except:
        text = None

    if not text:
        print(' ! Failed to get text:', fname, file=sys.stderr)
        return

    text = text.decode('utf8')
    text = re.sub('\n+', ' ', text)
    if len(text) < 800 or not ' the ' in text:
        print(' ! Not enough good text:', fname, file=sys.stderr)
        return

    d = Document()
    d.id = fname
    d.title = title
    d.url = url

    st = SentTokenizer()
    d.sections = [{'text': st.tokenize(text)}]

    return d


###

def chunks(l, n):
    """Yield successive n-sized chunks from l."""
    for i in range(0, len(l), n):
        yield l[i:i+n]

def ensure_directory(path):
    """Ensure that the desired path exists."""
    try:
        os.makedirs(path, exist_ok=True)
    except:
        raise

###

@click.command()
@click.option('--sd', is_flag=True, help='Use ScienceDirect.')
@click.option('--wiki', is_flag=True, help='Use Wikipedia.')
@click.option('--web', is_flag=True, help='Use tutorials from the Web.')
@click.argument('indir', type=click.Path(exists=True))
@click.argument('outdir', type=click.Path())
def main(indir, outdir, sd, wiki, web):
    global outputdir
    outputdir = outdir
    ensure_directory(outputdir + '/sd-download')
    ensure_directory(outputdir + '/web-download')

    expand = sd or wiki or web

    global pool
    pool = mp.Pool(PROCESSES)

    # If you're loading saved terms and want quicker results, change this to
    # c = Corpus(), but your output corpus won't include the original files.
    c = Corpus(indir, pool)

    terms = []
    if os.path.exists('data/terms.txt'):
        for line in open('data/terms.txt', 'r'):
            terms.append(line.strip())
        print('Loaded', len(terms), 'terms from data/terms.txt.')
    elif expand:
        terms = get_terms(c, 2000)
        print('Writing terms to data/terms.txt.')
        with open('data/terms.txt', 'w') as tout:
            for term in terms:
                tout.write(term + '\n')

    random.shuffle(terms)

    if len(terms) == 0 and expand:
        sys.stderr.write('Error: No terms to use.\n')
        sys.exit(1)

    # Download book chapters from ScienceDirect.
    if sd:
        docs_by_book = search_sd(terms)
        download_sd(docs_by_book)
        doc_ids = filter_sd(docs_by_book, terms)
        c |= get_sd_docs(doc_ids)

    # Download Wikipedia articles.
    if wiki:
        articles = search_wiki(terms)
        c |= get_wiki_docs(articles)

    # Download tutorial PDFs from the Web.
    if web:
        tutorials = search_web(terms)
        download_web(tutorials)
        c |= get_web_docs(tutorials)

    c.export(outputdir)


if __name__ == '__main__':
    pool = None
    outputdir = None
    main()