# This module is part of the Divmod project and is Copyright 2003 Amir Bakhtiar: # amir@divmod.org. This is free software; you can redistribute it and/or # modify it under the terms of version 2.1 of the GNU Lesser General Public # License as published by the Free Software Foundation. # import operator import string import math from sets import Set from splitter import Splitter class BayesData(dict): def __init__(self, name='', pool=None): self.name = name self.training = [] self.pool = pool self.tokenCount = 0 self.trainCount = 0 def trainedOn(self, item): return item in self.training def __repr__(self): return '<BayesDict: %s, %s tokens>' % (self.name, self.tokenCount) class Bayes(object): def __init__(self, tokenizer=None, combiner=None, dataClass=None): if dataClass is None: self.dataClass = BayesData else: self.dataClass = dataClass self.corpus = self.dataClass('__Corpus__') self.pools = {} self.pools['__Corpus__'] = self.corpus self.trainCount = 0 self.splitter = Splitter() self.dirty = True # The tokenizer takes an object and returns # a list of strings if tokenizer is None: self.tokenizer = self.getTokens else: self.tokenizer = tokenizer # The combiner combines probabilities if combiner is None: self.combiner = self.robinson else: self.combiner = combiner def split(self, text): return self.splitter.split(text) def commit(self): self.save() def newPool(self, poolName): """Create a new pool, without actually doing any training. """ self.dirty = True # not always true, but it's simple return self.pools.setdefault(poolName, self.dataClass(poolName)) def removePool(self, poolName): del(self.pools[poolName]) self.dirty = True def renamePool(self, poolName, newName): self.pools[newName] = self.pools[poolName] self.pools[newName].name = newName self.removePool(poolName) self.dirty = True def mergePools(self, destPool, sourcePool): """Merge an existing pool into another. The data from sourcePool is merged into destPool. The arguments are the names of the pools to be merged. The pool named sourcePool is left in tact and you may want to call removePool() to get rid of it. """ sp = self.pools[sourcePool] dp = self.pools[destPool] for tok, count in sp.items(): if dp.get(tok): dp[tok] += count else: dp[tok] = count dp.tokenCount += 1 self.dirty = True def poolData(self, poolName): """Return a list of the (token, count) tuples. """ return self.pools[poolName].items() def poolTokens(self, poolName): """Return a list of the tokens in this pool. """ return [tok for tok, count in self.poolData(poolName)] def save(self, fname='bayesdata.dat'): from cPickle import dump fp = open(fname, 'wb') dump(self.pools, fp) fp.close() def load(self, fname='bayesdata.dat'): from cPickle import load fp = open(fname, 'rb') self.pools = load(fp) fp.close() self.corpus = self.pools['__Corpus__'] self.dirty = True def poolNames(self): """Return a sorted list of Pool names. Does not include the system pool '__Corpus__'. """ pools = self.pools.keys() pools.remove('__Corpus__') pools = [pool for pool in pools] pools.sort() return pools def buildCache(self): """ merges corpora and computes probabilities """ self.cache = {} for pname, pool in self.pools.items(): # skip our special pool if pname == '__Corpus__': continue poolCount = len(pool) themCount = max(len(self.corpus) - poolCount, 1) cacheDict = self.cache.setdefault(pname, self.dataClass(pname)) for word, totCount in self.corpus.items(): # for every word in the copus # check to see if this pool contains this word thisCount = float(pool.get(word, 0.0)) otherCount = float(totCount) - thisCount if not poolCount: goodMetric = 1.0 else: goodMetric = min(1.0, otherCount/poolCount) badMetric = min(1.0, thisCount/themCount) f = badMetric / (goodMetric + badMetric) # PROBABILITY_THRESHOLD if abs(f-0.5) >= 0.1 : # GOOD_PROB, BAD_PROB cacheDict[word] = max(0.0001, min(0.9999, f)) def poolProbs(self): if self.dirty: self.buildCache() self.dirty = False return self.cache def getTokens(self, obj): """Hopefully it's a string and we'll just split it on non-alphanumeric stuff. Override this in your subclass for objects other than text. Alternatively, you can pass in a tokenizer as part of instance creation. """ return self.split(obj) def getProbs(self, pool, words): """ extracts the probabilities of tokens in a message """ probs = [(word, pool[word]) for word in words if word in pool] probs.sort(lambda x,y: cmp(y[1],x[1])) return probs[:2048] def train(self, pool, item, uid=None): """Train Bayes by telling him that item belongs in pool. uid is optional and may be used to uniquely identify the item that is being trained on. """ tokens = self.tokenizer(item) pool = self.pools.setdefault(pool, self.dataClass(pool)) self._train(pool, tokens) self.corpus.trainCount += 1 pool.trainCount += 1 if uid: pool.training.append(uid) self.dirty = True def untrain(self, pool, item, uid=None): tokens = self.tokenizer(item) pool = self.pools.get(pool, None) if not pool: return self._untrain(pool, tokens) # I guess we want to count this as additional training? self.corpus.trainCount += 1 pool.trainCount += 1 if uid: pool.training.remove(uid) self.dirty = True def _train(self, pool, tokens): wc = 0 for token in tokens: count = pool.get(token, 0) pool[token] = count + 1 count = self.corpus.get(token, 0) self.corpus[token] = count + 1 wc += 1 pool.tokenCount += wc self.corpus.tokenCount += wc def _untrain(self, pool, tokens): for token in tokens: count = pool.get(token, 0) if count: if count == 1: del(pool[token]) else: pool[token] = count - 1 pool.tokenCount -= 1 count = self.corpus.get(token, 0) if count: if count == 1: del(self.corpus[token]) else: self.corpus[token] = count - 1 self.corpus.tokenCount -= 1 def trainedOn(self, msg): for p in self.cache.values(): if msg in p.training: return True return False def guess(self, msg): tokens = Set(self.tokenizer(msg)) pools = self.poolProbs() res = {} for pname, pprobs in pools.items(): p = self.getProbs(pprobs, tokens) if len(p) != 0: res[pname]=self.combiner(p, pname) res = res.items() res.sort(lambda x,y: cmp(y[1], x[1])) return res def robinson(self, probs, ignore): """ computes the probability of a message being spam (Robinson's method) P = 1 - prod(1-p)^(1/n) Q = 1 - prod(p)^(1/n) S = (1 + (P-Q)/(P+Q)) / 2 Courtesy of http://christophe.delord.free.fr/en/index.html """ nth = 1./len(probs) P = 1.0 - reduce(operator.mul, map(lambda p: 1.0-p[1], probs), 1.0) ** nth Q = 1.0 - reduce(operator.mul, map(lambda p: p[1], probs)) ** nth S = (P - Q) / (P + Q) return (1 + S) / 2 def robinsonFisher(self, probs, ignore): """ computes the probability of a message being spam (Robinson-Fisher method) H = C-1( -2.ln(prod(p)), 2*n ) S = C-1( -2.ln(prod(1-p)), 2*n ) I = (1 + H - S) / 2 Courtesy of http://christophe.delord.free.fr/en/index.html """ n = len(probs) try: H = chi2P(-2.0 * math.log(reduce(operator.mul, map(lambda p: p[1], probs), 1.0)), 2*n) except OverflowError: H = 0.0 try: S = chi2P(-2.0 * math.log(reduce(operator.mul, map(lambda p: 1.0-p[1], probs), 1.0)), 2*n) except OverflowError: S = 0.0 return (1 + H - S) / 2 def __repr__(self): return '<Bayes: %s>' % [self.pools[p] for p in self.poolNames()] def __len__(self): return len(self.corpus) def chi2P(chi, df): """ return P(chisq >= chi, with df degree of freedom) df must be even """ assert df & 1 == 0 m = chi / 2.0 sum = term = math.exp(-m) for i in range(1, df/2): term *= m/i sum += term return min(sum, 1.0)