path: root/scripts/generic/balance-corpus

#!/usr/bin/ruby -w

require 'optparse'
require 'ostruct'
require 'pp'
require 'set'

options = OpenStruct.new
OptionParser.new { |opts|

	opts.banner = "Usage:	#{$0} [options]"

	opts.on("-n N","--num-parts N", Integer, "Number of parts into which the corpus should be split") { |v|
		options.parts = v
		options.parts_digits = options.parts.to_s.length
	}

	opts.on("-i FILE", "--corpus", String, "Corpus to split") { |v|
		options.corpus = v
	}

	options.reference = Array.new
	opts.on("-r FILE", "--reference", String, "Reference file") { |v|
		options.reference << v
	}
	
	options.put_all = false
	opts.on("-a","--all","Output all lines into a single file, in addition to split files") { |v|
		options.put_all = v
	}

	options.max_words = 1.0/0.0
	opts.on("-m N","--max-words", Integer, "Maximum number of words allowed in a line") { |v|
		options.max_words = v
	}

	options.min_words = 1
	opts.on("--min-words N", Integer, "Minimum number of words allowed in a line") { |v|
		options.min_words = v
	} 

	options.index_prefix = false
	opts.on("--index-prefix FILE_PREFIX", String, "Index file name prefixing the part number") { |v|
		options.index_prefix = v
	}

	opts.on("-p FILE_PREFIX","--prefix FILE_PREFIX", String, "File name prefixing the part number") { |v|
		options.output_prefix = v
	}

	opts.on("-s FILE_SUFFIX","--suffix FILE_SUFFIX", String, "File name suffixing the part number") { |v|
		options.output_suffix = v
	}

	options.ref_prefix = Array.new
	opts.on("--ref-prefix FILE_PREFIX", String, "File name prefixing the part number") { |v|
		options.ref_prefix << v
	}

	options.ref_suffix = Array.new
	opts.on("--ref-suffix FILE_SUFFIX", String, "File name suffixing the part number") { |v|
		options.ref_suffix << v
	}

	options.balance_naive = false
	opts.on("--balance-naive","Balance according to combined number of lines") { |v|
		options.balance_naive = v
	}

	options.balance_histogram = false
	opts.on("-h","--balance-histogram","Balance according to sentence length histogram") { |v|
		options.balance_histogram = v
	}

	options.balance_word_count = true
	opts.on("-w","--balance-words","Balance according to combined number of words") { |v|
		options.balance_word_count = v
	}

	options.balance_time = false
	opts.on("-t TIMES","--balance-time TIMES","Balance according to estimated per-sentence processing time") { |v|
		options.balance_time = v
	}

	options.verbose = false
	opts.on("-v","--[no-]verbose","Turn verbose on") { |v| 
		options.verbose = v
	}

	options.zero_pad = true
	opts.on("-z","--[no-]zeropad","Zero pad file names") { |v|
		options.zero_pad = v
	}

        if ARGV.length==0
	  puts opts
          exit
        end


}.parse!




class LineSize 
	include Comparable

	attr_reader :size, :index
	attr_writer :size

	@@max_index_digits = 0
	@@max_size_digits = 0
	
	def initialize(line,index)
		@index = index
		@size = line.strip.split(/\s+/).length

		index_digits = @index.to_s.length
		@@max_index_digits = index_digits if (index_digits > @@max_index_digits)

		size_digits = @size.to_s.length
		@@max_size_digits = size_digits if (size_digits > @@max_size_digits)
	end

	def <=>(other)
		if @size==other.size
			@index <=> other.index
		else
			size <=> other.size
		end
	end

	def to_s
		sprintf("Line %#{@@max_index_digits}i:	%#{@@max_size_digits}i words",@index, @size)
	end
end



def split_into_parts(file,part_for_line,parts,output_prefix,output_suffix,verbose,put_all,zeropad,index_prefix) 

	if (zeropad)
		parts_digits = parts.to_s.length
	else
		parts_digits = 0
	end

	out = Hash.new
	all = File.new("#{output_prefix}_all#{output_suffix}","a") if put_all
	index_out = Hash.new

	1.upto(parts) {|v|

		file_name = sprintf("%s%0#{parts_digits}i%s",output_prefix,v,output_suffix)
		out[v] = File.new(file_name,"w")
		
		unless index_prefix==false
			index_file_name = sprintf("%s%0#{parts_digits}i",index_prefix,v)
			index_out[v] = File.new(index_file_name,"w")
		end
	}


	File.open(file).each_with_index { |line,index|


		if (part_for_line.has_key?(index))
			puts "index==#{index}\tpart_for_line[#{index}]==#{part_for_line[index]}" if out[part_for_line[index]]==nil
			if verbose
				STDERR.puts "Line #{index} goes in #{out[part_for_line[index]].path}	#{line}"
			end	

			out[part_for_line[index]].puts(line)
			index_out[part_for_line[index]].puts(index) unless index_prefix==false

		elsif verbose
			STDERR.puts "Line #{index} will be skipped	#{line}"
		end
	}

	out.each_value { |file|
		file.close
	}


	if (put_all)
	    1.upto(parts) {|v|

		    file_name = sprintf("%s%0#{parts_digits}i%s",output_prefix,v,output_suffix)
		    File.open(file_name,"r").each { |line|
			    all.puts(line)
		    }

	    }

	    all.close
	end

end


def index_of_least(array)
	best=1.0/0 #Infinity
	best_index=0
	array.each_with_index {|v,i|
		if (v<best)
			best=v
			best_index=i
		end
	}
	return best_index
end


# Use to store which partition each line should be placed in
# 
# So, part_for_line[74] = 3 would mean that
# line number 74 should go into partition 3
#
part_for_line = Hash.new

# Use to store how many words are in each line
#
# So, words_per_line[74] = 15 would mean that
# line number 74 contains 15 words
#
words_per_line=Array.new

skipped_lines=Set.new

File.open(options.corpus).each_with_index { |line,index|

	line_size = LineSize.new(line,index)

	if line_size.size > options.max_words

		STDERR.puts "Line #{index} is too long: #{line_size.size} words. Max allowed is #{options.max_words}" if options.verbose
		skipped_lines.add(index)

	elsif line_size.size < options.min_words

		STDERR.puts "Line #{index} is too short: #{line_size.size} words. Min allowed is #{options.min_words}" if options.verbose
		skipped_lines.add(index)

	else

		words_per_line.push(line_size)

	end
}


if (options.balance_naive)

	total_lines=words_per_line.size

	STDERR.puts "total_lines=#{total_lines}" if options.verbose

	ceil=(total_lines/options.parts.to_f).ceil
	floor=(total_lines/options.parts.to_f).floor

	part_ceil = total_lines - floor*options.parts
	part_floor = options.parts - part_ceil

	STDERR.puts "#{ceil}*#{part_ceil} + #{floor}*#{part_floor}  =  #{ceil*part_ceil + floor*part_floor}" if options.verbose


	partition = 1
	lines_in_this_part = 0

	0.upto(total_lines-1) { |index|

		unless skipped_lines.include?(index)
			if (partition <= part_ceil)
				if (lines_in_this_part >= ceil)
					STDERR.puts "Partition #{partition} has #{lines_in_this_part} lines" if options.verbose
					lines_in_this_part=0 
					partition += 1
				end
			else
				if (lines_in_this_part >= floor)
					STDERR.puts "Partition #{partition} has #{lines_in_this_part} lines" if options.verbose
					lines_in_this_part=0
					partition += 1
				end
			end

			part_for_line[index] = partition
			lines_in_this_part += 1
			puts "part_for_line[#{index}] = #{partition}" if options.verbose
		end
		
	}

elsif (options.balance_histogram) 
    
    STDERR.puts "Balancing according to sentence length histogram"    
    
    words_per_line.sort!


    index=0

    words_per_line.each { |lineSize|
	    if index<options.parts
		    index+=1
	    else
		    index=1
	    end

	    part_for_line[lineSize.index] = index

    }

elsif (options.balance_word_count || options.balance_time)

	measure_unit = ""

	if (options.balance_time)
		STDERR.puts "Balancing according to time estimates"
		measure_unit = "seconds"

		index = 0
		File.open(options.balance_time).each_with_index { |time,line_index|
			unless (skipped_lines.include?(line_index))
				words_per_line[index].size = time.strip.to_f
				index += 1
			end
		}
				

	elsif (options.balance_word_count)
		STDERR.puts "Balancing according to word count"
		measure_unit = "words"
	end

	# Sort in reverse order
	words_per_line.sort! {|x,y| y <=> x }

	# Store the number of words that have been placed in each partition
	words_in_part = Array.new(options.parts,0)

	# At this point, words_per_line should be sorted with the longest sentences first
	words_per_line.each { |lineSize|
		partition = index_of_least(words_in_part)
		STDERR.puts "Line #{lineSize.index}\t#{lineSize.size} #{measure_unit}\tPartition #{partition}" if options.verbose
		part_for_line[lineSize.index] = partition+1 # part_for_line needs a 1-based partition index, so add 1
		words_in_part[partition] += lineSize.size
	}

	if (options.verbose)
		words_in_part.each_with_index { |words,partition|
			STDERR.puts "Partition #{partition}\t#{words} #{measure_unit}"
		}
	end

else


  exit;

end


split_into_parts(
	options.corpus,
	part_for_line,
	options.parts,
	options.output_prefix,
	options.output_suffix,
	options.verbose,
	options.put_all,
	options.zero_pad,
	options.index_prefix)



options.reference.each_with_index { |reference,index|

	split_into_parts(
		reference,
		part_for_line,
		options.parts,
		options.ref_prefix[index],
		options.ref_suffix[index],
		options.verbose,
		options.put_all,
		options.zero_pad,
		false)

}