<s> """<STR_LIT>""" <EOL> from __future__ import division , with_statement <EOL> from collections import defaultdict <EOL> from math import log <EOL> class ExpansionModel ( object ) : <EOL> def __init__ ( self , searcher , fieldname ) : <EOL> self . N = searcher . doc_count_all ( ) <EOL> self . collection_total = searcher . field_length ( fieldname ) <EOL> self . mean_length = self . collection_total / self . N <EOL> def normalizer ( self , maxweight , top_total ) : <EOL> raise NotImplementedError <EOL> def score ( self , weight_in_top , weight_in_collection , top_total ) : <EOL> raise NotImplementedError <EOL> class Bo1Model ( ExpansionModel ) : <EOL> def normalizer ( self , maxweight , top_total ) : <EOL> f = maxweight / self . N <EOL> return ( maxweight * log ( ( <NUM_LIT:1.0> + f ) / f ) + log ( <NUM_LIT:1.0> + f ) ) / log ( <NUM_LIT> ) <EOL> def score ( self , weight_in_top , weight_in_collection , top_total ) : <EOL> f = weight_in_collection / self . N <EOL> return weight_in_top * log ( ( <NUM_LIT:1.0> + f ) / f , <NUM_LIT:2> ) + log ( <NUM_LIT:1.0> + f , <NUM_LIT:2> ) <EOL> class Bo2Model ( ExpansionModel ) : <EOL> def normalizer ( self , maxweight , top_total ) : <EOL> f = maxweight * self . N / self . collection_total <EOL> return ( maxweight * log ( ( <NUM_LIT:1.0> + f ) / f , <NUM_LIT:2> ) + log ( <NUM_LIT:1.0> + f , <NUM_LIT:2> ) ) <EOL> def score ( self , weight_in_top , weight_in_collection , top_total ) : <EOL> f = weight_in_top * top_total / self . collection_total <EOL> return weight_in_top * log ( ( <NUM_LIT:1.0> + f ) / f , <NUM_LIT:2> ) + log ( <NUM_LIT:1.0> + f , <NUM_LIT:2> ) <EOL> class KLModel ( ExpansionModel ) : <EOL> def normalizer ( self , maxweight , top_total ) : <EOL> return maxweight * log ( self . collection_total / top_total ) / log ( <NUM_LIT> ) * top_total <EOL> def score ( self , weight_in_top , weight_in_collection , top_total ) : <EOL> wit_over_tt = weight_in_top / top_total <EOL> wic_over_ct = weight_in_collection / self . collection_total <EOL> if wit_over_tt < wic_over_ct : <EOL> return <NUM_LIT:0> <EOL> else : <EOL> return wit_over_tt * log ( ( wit_over_tt ) / ( weight_in_top / self . collection_total ) , <NUM_LIT:2> ) <EOL> class Expander ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , searcher , fieldname , model = Bo1Model ) : <EOL> """<STR_LIT>""" <EOL> self . fieldname = fieldname <EOL> if callable ( model ) : <EOL> model = model ( searcher , fieldname ) <EOL> self . model = model <EOL> term_reader = searcher . term_reader <EOL> self . collection_freq = dict ( ( word , freq ) for word , _ , freq <EOL> in term_reader . iter_field ( fieldname ) ) <EOL> self . topN_weight = defaultdict ( float ) <EOL> self . top_total = <NUM_LIT:0> <EOL> def add ( self , vector ) : <EOL> """<STR_LIT>""" <EOL> total_weight = <NUM_LIT:0> <EOL> topN_weight = self . topN_weight <EOL> for word , weight in vector : <EOL> total_weight += weight <EOL> topN_weight [ word ] += weight <EOL> self . top_total += total_weight <EOL> def expanded_terms ( self , number , normalize = True ) : <EOL> """<STR_LIT>""" <EOL> model = self . model <EOL> tlist = [ ] <EOL> maxweight = <NUM_LIT:0> <EOL> collection_freq = self . collection_freq <EOL> for word , weight in self . topN_weight . iteritems ( ) : <EOL> score = model . score ( weight , collection_freq [ word ] , self . top_total ) <EOL> if score > maxweight : maxweight = score <EOL> tlist . append ( ( score , word ) ) <EOL> if normalize : <EOL> norm = model . normalizer ( maxweight , self . top_total ) <EOL> else : <EOL> norm = maxweight <EOL> tlist = [ ( weight / norm , t ) for weight , t in tlist ] <EOL> tlist . sort ( reverse = True ) <EOL> return [ ( t , weight ) for weight , t in tlist [ : number ] ] </s>
