<s> """<STR_LIT>""" <EOL> from collections import OrderedDict <EOL> from six . moves import range <EOL> import numpy <EOL> from picklable_itertools . extras import equizip <EOL> from theano import config , function <EOL> from blocks_extras . bricks . sequence_generator2 import SequenceGenerator <EOL> from blocks . filter import VariableFilter , get_application_call , get_brick <EOL> from blocks . graph import ComputationGraph <EOL> from blocks . roles import INPUT , OUTPUT <EOL> from blocks . utils import unpack <EOL> class BeamSearch ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , samples ) : <EOL> self . cg = ComputationGraph ( samples ) <EOL> self . inputs = self . cg . inputs <EOL> self . generator = get_brick ( samples ) <EOL> if not isinstance ( self . generator , SequenceGenerator ) : <EOL> raise ValueError <EOL> self . generate_call = get_application_call ( samples ) <EOL> if ( not self . generate_call . application == <EOL> self . generator . generate ) : <EOL> raise ValueError <EOL> self . inner_cg = ComputationGraph ( self . generate_call . inner_outputs ) <EOL> self . context_names = self . generator . generate . contexts <EOL> self . state_names = self . generator . generate . states <EOL> self . contexts = [ <EOL> VariableFilter ( bricks = [ self . generator ] , <EOL> name = name , <EOL> roles = [ INPUT ] ) ( self . inner_cg ) [ <NUM_LIT:0> ] <EOL> for name in self . context_names ] <EOL> self . input_states = [ ] <EOL> self . input_state_names = [ ] <EOL> for name in self . generator . generate . states : <EOL> var = VariableFilter ( <EOL> bricks = [ self . generator ] , name = name , <EOL> roles = [ INPUT ] ) ( self . inner_cg ) <EOL> if var : <EOL> self . input_state_names . append ( name ) <EOL> self . input_states . append ( var [ <NUM_LIT:0> ] ) <EOL> self . compiled = False <EOL> def _compile_initial_state_and_context_computer ( self ) : <EOL> initial_states = VariableFilter ( <EOL> applications = [ self . generator . initial_states ] , <EOL> roles = [ OUTPUT ] ) ( self . cg ) <EOL> outputs = OrderedDict ( [ ( v . tag . name , v ) for v in initial_states ] ) <EOL> beam_size = unpack ( VariableFilter ( <EOL> applications = [ self . generator . initial_states ] , <EOL> name = '<STR_LIT>' ) ( self . cg ) ) <EOL> for name , context in equizip ( self . context_names , self . contexts ) : <EOL> outputs [ name ] = context <EOL> outputs [ '<STR_LIT>' ] = beam_size <EOL> self . initial_state_and_context_computer = function ( <EOL> self . inputs , outputs , on_unused_input = '<STR_LIT:ignore>' ) <EOL> def _compile_next_state_computer ( self ) : <EOL> next_states = [ VariableFilter ( bricks = [ self . generator ] , <EOL> name = name , <EOL> roles = [ OUTPUT ] ) ( self . inner_cg ) [ - <NUM_LIT:1> ] <EOL> for name in self . state_names ] <EOL> samples = VariableFilter ( <EOL> applications = [ self . generator . readout . sample ] , name = '<STR_LIT>' ) ( <EOL> self . inner_cg . variables ) <EOL> self . next_state_computer = function ( <EOL> self . contexts + self . input_states + samples , next_states ) <EOL> def _compile_scores_computer ( self ) : <EOL> scores = self . generator . readout . scores ( <EOL> ** dict ( zip ( self . input_state_names , self . input_states ) + <EOL> zip ( self . context_names , self . contexts ) ) ) <EOL> self . scores_computer = function ( <EOL> self . contexts + self . input_states , scores , <EOL> on_unused_input = '<STR_LIT:ignore>' ) <EOL> def compile ( self ) : <EOL> """<STR_LIT>""" <EOL> self . _compile_initial_state_and_context_computer ( ) <EOL> self . _compile_next_state_computer ( ) <EOL> self . _compile_scores_computer ( ) <EOL> self . compiled = True <EOL> def compute_initial_states_and_contexts ( self , inputs ) : <EOL> """<STR_LIT>""" <EOL> outputs = self . initial_state_and_context_computer ( <EOL> * [ inputs [ var ] for var in self . inputs ] ) <EOL> contexts = OrderedDict ( ( n , outputs . pop ( n ) ) for n in self . context_names ) <EOL> beam_size = outputs . pop ( '<STR_LIT>' ) <EOL> initial_states = outputs <EOL> return contexts , initial_states , beam_size <EOL> def compute_scores ( self , contexts , states ) : <EOL> """<STR_LIT>""" <EOL> input_states = [ states [ name ] for name in self . input_state_names ] <EOL> return self . scores_computer ( * ( list ( contexts . values ( ) ) + <EOL> input_states ) ) <EOL> def compute_next_states ( self , contexts , states , outputs ) : <EOL> """<STR_LIT>""" <EOL> input_states = [ states [ name ] for name in self . input_state_names ] <EOL> next_values = self . next_state_computer ( * ( list ( contexts . values ( ) ) + <EOL> input_states + [ outputs ] ) ) <EOL> return OrderedDict ( equizip ( self . state_names , next_values ) ) <EOL> @ staticmethod <EOL> def _smallest ( matrix , k , only_first_row = False ) : <EOL> """<STR_LIT>""" <EOL> if only_first_row : <EOL> flatten = matrix [ : <NUM_LIT:1> , : ] . flatten ( ) <EOL> else : <EOL> flatten = matrix . flatten ( ) <EOL> args = numpy . argpartition ( flatten , k ) [ : k ] <EOL> args = args [ numpy . argsort ( flatten [ args ] ) ] <EOL> return numpy . unravel_index ( args , matrix . shape ) , flatten [ args ] <EOL> def search ( self , input_values , eol_symbol , max_length , <EOL> ignore_first_eol = False , as_arrays = False ) : <EOL> """<STR_LIT>""" <EOL> if not self . compiled : <EOL> self . compile ( ) <EOL> contexts , states , beam_size = self . compute_initial_states_and_contexts ( <EOL> input_values ) <EOL> all_outputs = - <NUM_LIT:1> * numpy . ones ( ( <NUM_LIT:1> , beam_size ) , dtype = '<STR_LIT>' ) <EOL> all_masks = numpy . ones_like ( all_outputs , dtype = config . floatX ) <EOL> all_costs = numpy . zeros_like ( all_outputs , dtype = config . floatX ) <EOL> for i in range ( max_length ) : <EOL> if all_masks [ - <NUM_LIT:1> ] . sum ( ) == <NUM_LIT:0> : <EOL> break <EOL> scores = self . compute_scores ( contexts , states ) <EOL> next_costs = ( all_costs [ - <NUM_LIT:1> , : , None ] - <EOL> scores * all_masks [ - <NUM_LIT:1> , : , None ] ) <EOL> ( finished , ) = numpy . where ( all_masks [ - <NUM_LIT:1> ] == <NUM_LIT:0> ) <EOL> next_costs [ finished , : eol_symbol ] = numpy . inf <EOL> next_costs [ finished , eol_symbol + <NUM_LIT:1> : ] = numpy . inf <EOL> ( indexes , outputs ) , chosen_costs = self . _smallest ( <EOL> next_costs , beam_size , only_first_row = i == <NUM_LIT:0> ) <EOL> for name in states : <EOL> states [ name ] = states [ name ] [ indexes ] <EOL> all_outputs = all_outputs [ : , indexes ] <EOL> all_masks = all_masks [ : , indexes ] <EOL> all_costs = all_costs [ : , indexes ] <EOL> states . update ( self . compute_next_states ( contexts , states , outputs ) ) <EOL> all_outputs = numpy . vstack ( [ all_outputs , outputs [ None , : ] ] ) <EOL> all_costs = numpy . vstack ( [ all_costs , chosen_costs [ None , : ] ] ) <EOL> mask = outputs != eol_symbol <EOL> if ignore_first_eol and i == <NUM_LIT:0> : <EOL> mask [ : ] = <NUM_LIT:1> <EOL> all_masks = numpy . vstack ( [ all_masks , mask [ None , : ] ] ) <EOL> all_outputs = all_outputs [ <NUM_LIT:1> : ] <EOL> all_masks = all_masks [ : - <NUM_LIT:1> ] <EOL> all_costs = all_costs [ <NUM_LIT:1> : ] - all_costs [ : - <NUM_LIT:1> ] <EOL> result = all_outputs , all_masks , all_costs <EOL> if as_arrays : <EOL> return result <EOL> return self . result_to_lists ( result ) <EOL> @ staticmethod <EOL> def result_to_lists ( result ) : <EOL> outputs , masks , costs = [ array . T for array in result ] <EOL> outputs = [ list ( output [ : mask . sum ( ) ] ) <EOL> for output , mask in equizip ( outputs , masks ) ] <EOL> costs = list ( costs . T . sum ( axis = <NUM_LIT:0> ) ) <EOL> return outputs , costs </s>
