<s> import sqlite3 <EOL> import operator <EOL> import random <EOL> from csv import DictWriter <EOL> from collections import defaultdict <EOL> from extract_features import instantiate_feature , guesses_for_question <EOL> from util . qdb import QuestionDatabase <EOL> from unidecode import unidecode <EOL> kCOUNT_CUTOFF = <NUM_LIT:2> <EOL> kCHOICEIDS = "<STR_LIT>" <EOL> kCATEGORIES = set ( """<STR_LIT>""" . split ( "<STR_LIT:\n>" ) ) <EOL> class McScience : <EOL> def __init__ ( self , page , question , fold ) : <EOL> self . page = page <EOL> self . question = question <EOL> self . choices = [ ] <EOL> self . fold = fold <EOL> def add_text ( self , text ) : <EOL> self . text = text <EOL> def add_choices ( self , choices ) : <EOL> self . choices = list ( choices ) <EOL> random . shuffle ( self . choices ) <EOL> def csv_line ( self , choice_strings , destination = "<STR_LIT:train>" ) : <EOL> d = { } <EOL> d [ "<STR_LIT:id>" ] = self . question <EOL> if destination != "<STR_LIT:key>" : <EOL> d [ "<STR_LIT>" ] = self . text <EOL> for ii , cc in enumerate ( self . choices ) : <EOL> if destination != "<STR_LIT:key>" : <EOL> d [ "<STR_LIT>" % choice_strings [ ii ] ] = unidecode ( cc ) <EOL> if cc == self . page and ( destination == "<STR_LIT:train>" or destination == "<STR_LIT:key>" ) : <EOL> d [ "<STR_LIT>" ] = choice_strings [ ii ] <EOL> assert self . page in self . choices , "<STR_LIT>" % ( self . page , str ( self . choices ) ) <EOL> return d <EOL> def question_top_guesses ( text , deep , guess_connection , id , page , num_guesses = <NUM_LIT:4> ) : <EOL> """<STR_LIT>""" <EOL> c = guess_connection . cursor ( ) <EOL> command = ( '<STR_LIT>' + <EOL> '<STR_LIT>' ) % ( id , num_guesses + <NUM_LIT:1> ) <EOL> c . execute ( command ) <EOL> choices = set ( [ page ] ) <EOL> for ii , in c : <EOL> if len ( choices ) < num_guesses and not ii in choices : <EOL> choices . add ( ii ) <EOL> new_guesses = deep . text_guess ( text ) <EOL> for guess , score in sorted ( new_guesses . items ( ) , key = operator . itemgetter ( <NUM_LIT:1> ) , reverse = True ) : <EOL> if len ( choices ) < num_guesses and not guess in choices : <EOL> choices . add ( guess ) <EOL> return choices <EOL> def question_first_sentence ( database_connection , question ) : <EOL> """<STR_LIT>""" <EOL> c = database_connection . cursor ( ) <EOL> command = '<STR_LIT>' % question <EOL> c . execute ( command ) <EOL> for ii , in c : <EOL> return unidecode ( ii ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> import argparse <EOL> parser = argparse . ArgumentParser ( description = '<STR_LIT>' ) <EOL> default_path = '<STR_LIT>' <EOL> parser . add_argument ( '<STR_LIT>' , type = str , default = default_path + '<STR_LIT>' ) <EOL> parser . add_argument ( '<STR_LIT>' , type = str , default = default_path + '<STR_LIT>' , <EOL> help = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , type = int , default = <NUM_LIT:4> , <EOL> help = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , type = str , default = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , type = str , default = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , type = str , default = "<STR_LIT>" ) <EOL> flags = parser . parse_args ( ) <EOL> print ( "<STR_LIT>" % flags . question_db ) <EOL> question_database = sqlite3 . connect ( flags . question_db ) <EOL> guess_database = sqlite3 . connect ( flags . guess_db ) <EOL> query = '<STR_LIT>' <EOL> query += "<STR_LIT>" . join ( "<STR_LIT>" % x for x in kCATEGORIES ) <EOL> query += "<STR_LIT:)>" <EOL> c = question_database . cursor ( ) <EOL> print ( query ) <EOL> c . execute ( query ) <EOL> answer_count = defaultdict ( int ) <EOL> for pp , in c : <EOL> answer_count [ pp ] += <NUM_LIT:1> <EOL> query = '<STR_LIT>' <EOL> c = question_database . cursor ( ) <EOL> c . execute ( query ) <EOL> print ( list ( x for x in answer_count if answer_count [ x ] >= kCOUNT_CUTOFF ) ) <EOL> print ( len ( list ( x for x in answer_count if answer_count [ x ] >= kCOUNT_CUTOFF ) ) ) <EOL> deep = instantiate_feature ( "<STR_LIT>" , QuestionDatabase ( flags . question_db ) ) <EOL> questions = { } <EOL> question_num = <NUM_LIT:0> <EOL> for pp , ii , nn , ff in c : <EOL> if nn >= <NUM_LIT:0> or answer_count [ pp ] < kCOUNT_CUTOFF : <EOL> continue <EOL> question_num += <NUM_LIT:1> <EOL> question = McScience ( pp , ii , ff ) <EOL> question . add_text ( question_first_sentence ( question_database , ii ) ) <EOL> choices = question_top_guesses ( question . text , deep , guess_database , ii , pp , <EOL> flags . num_choices ) <EOL> question . add_choices ( choices ) <EOL> questions [ ii ] = question <EOL> if question_num % <NUM_LIT:100> == <NUM_LIT:0> : <EOL> print ( pp , ii , question_num ) <EOL> print ( choices ) <EOL> answer_choices = [ "<STR_LIT>" % kCHOICEIDS [ x ] for x in xrange ( flags . num_choices ) ] <EOL> train_out = DictWriter ( open ( flags . train_out , '<STR_LIT:w>' ) , [ "<STR_LIT:id>" , "<STR_LIT>" , "<STR_LIT>" ] + <EOL> answer_choices ) <EOL> train_out . writeheader ( ) <EOL> test_out = DictWriter ( open ( flags . test_out , '<STR_LIT:w>' ) , [ "<STR_LIT:id>" , "<STR_LIT>" ] + answer_choices ) <EOL> test_out . writeheader ( ) <EOL> key_out = DictWriter ( open ( flags . key_out , '<STR_LIT:w>' ) , [ "<STR_LIT:id>" , "<STR_LIT>" ] ) <EOL> key_out . writeheader ( ) <EOL> for qq in questions . values ( ) : <EOL> print ( qq . fold ) <EOL> if qq . fold == "<STR_LIT>" : <EOL> test_out . writerow ( qq . csv_line ( kCHOICEIDS , "<STR_LIT:test>" ) ) <EOL> key_out . writerow ( qq . csv_line ( kCHOICEIDS , "<STR_LIT:key>" ) ) <EOL> else : <EOL> train_out . writerow ( qq . csv_line ( kCHOICEIDS , "<STR_LIT:train>" ) ) </s>
