<s> from __future__ import division <EOL> from collections import Counter , defaultdict <EOL> from json import loads , dumps <EOL> from stat_parser . word_classes import word_class <EOL> class PCFG : <EOL> RARE_WORD_COUNT = <NUM_LIT:5> <EOL> def __init__ ( self ) : <EOL> self . q1 = defaultdict ( float ) <EOL> self . q2 = defaultdict ( float ) <EOL> self . well_known_words = set ( ) <EOL> def norm_word ( self , word ) : <EOL> return word if word in self . well_known_words else word_class ( word ) <EOL> def __build_caches ( self ) : <EOL> self . N = set ( ) <EOL> self . binary_rules = defaultdict ( list ) <EOL> for x , _ in self . q1 . keys ( ) : <EOL> self . N . add ( x ) <EOL> for x , y1 , y2 in self . q2 . keys ( ) : <EOL> self . N . update ( set ( [ x , y1 , y2 ] ) ) <EOL> self . binary_rules [ x ] . append ( ( y1 , y2 ) ) <EOL> def learn_from_treebanks ( self , treebanks ) : <EOL> self . sym_count = Counter ( ) <EOL> self . unary_count = Counter ( ) <EOL> self . binary_count = Counter ( ) <EOL> self . words_count = Counter ( ) <EOL> for treebank in treebanks : <EOL> for s in open ( treebank ) : <EOL> self . __count ( loads ( s ) ) <EOL> for word , count in self . words_count . iteritems ( ) : <EOL> if count >= PCFG . RARE_WORD_COUNT : <EOL> self . well_known_words . add ( word ) <EOL> norm = Counter ( ) <EOL> for ( x , word ) , count in self . unary_count . iteritems ( ) : <EOL> norm [ ( x , self . norm_word ( word ) ) ] += count <EOL> self . unary_count = norm <EOL> for ( x , word ) , count in self . unary_count . iteritems ( ) : <EOL> self . q1 [ x , word ] = self . unary_count [ x , word ] / self . sym_count [ x ] <EOL> for ( x , y1 , y2 ) , count in self . binary_count . iteritems ( ) : <EOL> self . q2 [ x , y1 , y2 ] = self . binary_count [ x , y1 , y2 ] / self . sym_count [ x ] <EOL> self . __build_caches ( ) <EOL> def __count ( self , tree ) : <EOL> if isinstance ( tree , basestring ) : return <EOL> sym = tree [ <NUM_LIT:0> ] <EOL> self . sym_count [ sym ] += <NUM_LIT:1> <EOL> if len ( tree ) == <NUM_LIT:3> : <EOL> y1 , y2 = ( tree [ <NUM_LIT:1> ] [ <NUM_LIT:0> ] , tree [ <NUM_LIT:2> ] [ <NUM_LIT:0> ] ) <EOL> self . binary_count [ ( sym , y1 , y2 ) ] += <NUM_LIT:1> <EOL> self . __count ( tree [ <NUM_LIT:1> ] ) <EOL> self . __count ( tree [ <NUM_LIT:2> ] ) <EOL> elif len ( tree ) == <NUM_LIT:2> : <EOL> word = tree [ <NUM_LIT:1> ] <EOL> self . unary_count [ ( sym , word ) ] += <NUM_LIT:1> <EOL> self . words_count [ word ] += <NUM_LIT:1> <EOL> def save_model ( self , path ) : <EOL> with open ( path , '<STR_LIT:w>' ) as model : <EOL> for ( x , word ) , p in self . q1 . iteritems ( ) : <EOL> model . write ( dumps ( [ '<STR_LIT>' , x , word , p ] ) + '<STR_LIT:\n>' ) <EOL> for ( x , y1 , y2 ) , p in self . q2 . iteritems ( ) : <EOL> model . write ( dumps ( [ '<STR_LIT>' , x , y1 , y2 , p ] ) + '<STR_LIT:\n>' ) <EOL> model . write ( dumps ( [ '<STR_LIT>' , list ( self . well_known_words ) ] ) + '<STR_LIT:\n>' ) <EOL> def load_model ( self , path ) : <EOL> with open ( path ) as model : <EOL> for line in model : <EOL> data = loads ( line ) <EOL> if data [ <NUM_LIT:0> ] == '<STR_LIT>' : <EOL> _ , x , word , p = data <EOL> self . q1 [ x , word ] = p <EOL> elif data [ <NUM_LIT:0> ] == '<STR_LIT>' : <EOL> _ , x , y1 , y2 , p = data <EOL> self . q2 [ x , y1 , y2 ] = p <EOL> elif data [ <NUM_LIT:0> ] == '<STR_LIT>' : <EOL> self . well_known_words = data [ <NUM_LIT:1> ] <EOL> self . __build_caches ( ) </s>
