gramificate/gramificate.py

#!/usr/bin/python

import sys
import re

class NGramSet:
	def __init__(self, filename, gram_size):
		self.gram_size = gram_size
		self.filename = filename
		self.text = None
		self.gram_set = {}
		self.total_count = 0

	def read_file(self):
		# open and read file
		src = open(self.filename)
		text = src.read()

		# normalize text to remove extra space and characters
		text = text.lower() # lower case chars
		text = re.sub('[^a-z]', ' ', text) # remove anything not a-z
		text = re.sub('  *', ' ', text) # shrink all multi spaces to a single
		self.text= text


	# takes in text and populates class with ngram info
	def process(self):
		if (self.text == None):
			self.read_file()
		gram = []
		for word in self.text.split(" "):
			if (word == ''):
				continue
			if (len(gram) >= self.gram_size):
				gram.pop(0)
			gram.append(word)
			if (len(gram) == self.gram_size):
				self.total_count += 1
				key = ' '.join(gram)
				if (not self.gram_set.has_key(key)):
					self.gram_set[key] = {'gram': list(gram), 'count': 1}
				else:
					self.gram_set[key]['count'] += 1
		self.generate_gram_stats()

	def generate_gram_stats(self):
		for key in self.gram_set:
			self.gram_set[key]['percent'] = self.gram_set[key]['count'] / float(self.total_count)

	def print_mysql(self):
		print "mysql"

if len(sys.argv) < 3:
	print "Usage: gramificate.py [N] [FILE]"
	print "  N     Gram size"
	print "  FILE  Filename"
	print "Output: Mysql commands to create a table FILE-N that contains all the grams and" 
	print "        associated stats (count of gram, percent of total). Can be directly and"
	print "        safely piped into mysql:"
	print "          mysql -u USER -pPass -D gram_db < ./gramificate.py 2 input.txt" 
	exit()

# generate ngrams
grams = NGramSet(sys.argv[2], int(sys.argv[1]))
grams.process()
print grams.gram_set