bakeoff-sample/run_textrank_ver1.py

import pandas as pd
import argparse
import os
import re
import csv
import spacy
import json
import numpy as np
import logging

import traceback

import dagshub

from summa.summarizer import summarize
from gensim.utils import simple_preprocess
from gensim.parsing.preprocessing import remove_stopwords

from rouge import Rouge
from nltk.corpus import stopwords

stopwords = set(stopwords.words('english'))
nlp = spacy.load('en_core_web_sm')

cache = {}

class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            return super(NpEncoder, self).default(obj)


def clean_sentence(text):
    """ Clean an input sentence
    Args:
        text (str): the sentence
    Returns:
        text (str): the cleaned sentence
    """
    sent = nlp(text)
    # clean the first few characters
    result = text
    for i, token in enumerate(sent):
        if token.pos_ not in ['PUNCT', 'X', 'SYM']:
            result = sent[i:].text
            break

    # capitalize the first character
    if len(result) > 0:
        result = result[0].upper() + result[1:]

    return result

def drop_unwanted_sentences(answers):
    """ Remove unwanted sentences and concatenate into one string
    Args:
        answers (list of str): the answers to be cleaned
    Returns:
        str: the cleaned and concatenated string
    """
    all_sents = []
    for text in answers:
        sents = nlp(text, disable=['tagger', 'ner']).sents
        for sent in sents:
            if len(sent) <= 30:
                all_sents.append(sent.text)

    return '\n'.join(all_sents)

def do_textrank(qid, num_sent=5, num_words=100):
    """ Perform TextRank to generate the summary items with coverage information
    Args:
        qid (str): the question_id to be summarized
        num_sent (int, optional): the maximal number of sentences in the summary
        num_words (int, optional): the maximal number of words in the summary
    Returns:
        Dict: the summary result containing the summary items,
            coverage info, and related answers.
    """
    answers = df.loc[df['questionId'] == qid]
    if len(answers) <= num_sent:
        return None


    company_id = answers['fccompanyId'].iloc[0]
    company_name = answers['companyName'].iloc[0].strip()
    answer_ids = list(answers['answerId'])
    answer_text = list(answers['content_answer'])
    answer_dict = dict(zip(answer_ids, answer_text))
    question_text = answers['content_question'].iloc[0]
    question_code = answers['questionCode'].iloc[0]
    if answers['questionCode'].isna().iloc[0]:
        question_code = ''
    question_topics = answers['topics'].iloc[0]
    result = {'question_id': qid,
              'question_text': question_text,
              'question_code': question_code,
              'question_topics': question_topics,
              'company_id': company_id,
              'company': company_name,
              'num_answers': len(answers)}

    # print(result)
    # input('press any key to continue')
    logging.debug("Input:" + str(answer_text))
    full_text = drop_unwanted_sentences(answer_text)
    logging.debug("Text:" + full_text)
    ans = summarize(full_text, words=num_words * 2)
    summary = []
    summary_lines = []

    stats = {}
    stats['num_summaries'] = 0
    stats['num_answers'] = len(answers)


    for idx, line in enumerate(ans.split('\n')[:num_sent]):
        summary_lines.append(clean_sentence(line))

    stats['num_summaries'] = len(summary_lines)
   
    for idx, line in enumerate(summary_lines):

        summary_item = {}
        summary_item['text'] = line

        summary.append(summary_item)


    result['summary'] = summary


    rouge = Rouge()
    scores = rouge.get_scores(ans, full_text)[0]

    stats['ROUGE-1'] = scores['rouge-1']['f']
    stats['ROUGE-2'] = scores['rouge-2']['f']
    stats['ROUGE-L'] = scores['rouge-l']['f']
    result['stats'] = stats

    logging.debug(result)
    return result


def find_questions(keywords, topics, qids):
    c = questions['content'].str.contains('')
    if len(keywords) > 0:
        ck = questions['content'].str.contains('')
        for keyword in keywords:
            ck = ck & questions['content'].str.contains(keyword)
        c = c & ck

    if len(topics) > 0:
        ct = questions['topics'].str.contains(topics[0]) | questions['questionCode'].str.contains(topics[0])
        for topic in topics[1:]:
            ct = ct | ( questions['topics'].str.contains(topic) | questions['questionCode'].str.contains(topic) )
        c = c & ct

    if len(qids) > 0:
        qt = questions['questionId'].str.contains('')
        for qid in qids:
            qt = qt & questions['questionId'].str.contains(qid)
        c = c & qt

    return questions[c]

def write_json_to_file(data, file_name):
    with open(file_name, 'w') as fo:
        json.dump(data, fo, cls=NpEncoder)


def summarize_with_textrank(questions,
                            num_sent=5,
                            num_words=100):

    results = []

    stats = {}
    stats['question_count'] = 0
    stats['summarized_question_count'] = 0
    stats['total_num_summaries'] = 0
    stats['total_num_answers'] = 0
    stats['avg_ROUGE-1'] = 0
    stats['avg_ROUGE-2'] = 0
    stats['avg_ROUGE-L'] = 0


    for qid in questions['questionId']:

        if stats['question_count'] % 100 == 0:
            logging.info("Processed {} questions summarized {}".format(stats['question_count'],stats['summarized_question_count']))

        stats['question_count'] = stats['question_count'] + 1
        try:
            result = do_textrank(qid, num_sent, num_words)


            if result != None:
                stats['avg_ROUGE-1'] = ( stats['avg_ROUGE-1'] * stats['summarized_question_count'] + result['stats']['ROUGE-1'] ) / ( stats['summarized_question_count'] + 1)
                stats['avg_ROUGE-2'] = ( stats['avg_ROUGE-2'] * stats['summarized_question_count'] + result['stats']['ROUGE-2'] ) / ( stats['summarized_question_count'] + 1)
                stats['avg_ROUGE-L'] = ( stats['avg_ROUGE-L'] * stats['summarized_question_count'] + result['stats']['ROUGE-L'] ) / ( stats['summarized_question_count'] + 1)
                stats['summarized_question_count'] = stats['summarized_question_count'] + 1
                stats['total_num_summaries'] = stats['total_num_summaries'] + result['stats']['num_summaries']
                stats['total_num_answers'] =  stats['total_num_answers'] + result['stats']['num_answers']
                
                results.append(result)

                #with dagshub.dagshub_logger() as logger:
                logger.log_metrics(stats)
                
                write_json_to_file(result, "data/output/" + qid + ".json")
            else:
                continue
                # write_json_to_file({"message": "no summary created"}, "../data/output/" + qid + ".json")


        except Exception as exception:
            traceback.print_exc()
            logging.error("Exception on question: " + qid)

    logging.info("Final statistics {}".format(stats))

if __name__ == '__main__':
    # python run_textrank.py --keywords interview,precefiss --num_sentences 5 --num_words 100
    parser = argparse.ArgumentParser()
    parser.add_argument("--keywords", type=str, default=None)
    parser.add_argument("--num_sentences", type=int, default=5)
    parser.add_argument("--num_words", type=int, default=100)
    parser.add_argument("--topics", type=str, default=None)
    parser.add_argument("--questions", type=str, default=None)
    parser.add_argument("--log", type=str, default=None)
    hp = parser.parse_args()

    # Initialize logging
    logging_level = logging.INFO if hp.log is None else hp.log.upper()
    logging.basicConfig(format='%(levelname)s:%(message)s', level=logging_level)

    # read csv into dataframes
    questions = pd.read_csv('data/questions.csv', warn_bad_lines=True, error_bad_lines=False, verbose=True)
    answers = pd.read_csv('data/answers.csv', warn_bad_lines=True, error_bad_lines=False, verbose=True)
    companies = pd.read_csv('data/fccid-companyName.csv', warn_bad_lines=True, error_bad_lines=False, verbose=True)
    df = pd.merge(answers, questions, on='questionId', suffixes=('_answer','_question'))
    df = pd.merge(df, companies, on='fccompanyId')

    # filter questions based on keywords, topics, and/or quesstion ids
    keywords = [] if hp.keywords is None else hp.keywords.split(',')
    topics = [] if hp.topics is None else hp.topics.split(',')
    qids = [] if hp.questions is None else hp.questions.split(',')
    filtered_questions = find_questions(keywords, topics, qids)

    print(vars(hp))

    with dagshub.dagshub_logger() as logger:
        logger.log_hyperparams(vars(hp))

        # summarize
        summarize_with_textrank(filtered_questions,
                            num_sent=hp.num_sentences,
                            num_words=hp.num_words)