TOP

mailnews/extensions/bayesian-spam-filter/src/nsBayesianFilter.cpp
 Token* Tokenizer::add(const char* word, PRUint32 count)
 {
     PLDHashEntryHdr* entry = PL_DHashTableOperate(&mTokenTable, word, PL_DHASH_ADD);
     Token* token = NS_STATIC_CAST(Token*, entry);
     if (token) {
         if (token->mWord == NULL) {
             PRUint32 len = strlen(word);
             NS_ASSERTION(len != 0, "adding zero length word to tokenizer");
             token->mWord = copyWord(word, len);
             NS_ASSERTION(token->mWord, "copyWord failed");
             if (!token->mWord) {
                 PL_DHashTableRawRemove(&mTokenTable, entry);
                 return NULL;
             }
             token->mLength = len;
             token->mCount = count;
             token->mProbability = 0;
         } else {
             token->mCount += count;
         }
     }
     return token;
 }

 void nsBayesianFilter::classifyMessage(Tokenizer& tokenizer, const char* messageURI,
                                        nsIJunkMailClassificationListener* listener)
 {
     Token* tokens = tokenizer.copyTokens();
     if (!tokens) return;
   
     // the algorithm in "A Plan For Spam" assumes that you have a large good
     // corpus and a large junk corpus.
     // that won't be the case with users who first use the junk mail feature
     // so, we do certain things to encourage them to train.
     //
     // if there are no good tokens, assume the message is junk
     // this will "encourage" the user to train
     // and if there are no bad tokens, assume the message is not junk
     // this will also "encourage" the user to train
     // see bug #194238
     if (listener && !mGoodCount && !mGoodTokens.countTokens()) {
       listener->OnMessageClassified(messageURI, nsMsgJunkStatus(nsIJunkMailPlugin::JUNK));
       return;
     }
     else if (listener && !mBadCount && !mBadTokens.countTokens()) {
       listener->OnMessageClassified(messageURI, nsMsgJunkStatus(nsIJunkMailPlugin::GOOD));
       return;
     }
 
     /* run the kernel of the Graham filter algorithm here. */
     PRUint32 i, count = tokenizer.countTokens();
     double ngood = mGoodCount, nbad = mBadCount;
     for (i = 0; i < count; ++i) {
         Token& token = tokens[i];
         const char* word = token.mWord;
         // ((g (* 2 (or (gethash word good) 0)))
         Token* t = mGoodTokens.get(word);
         double g = 2.0 * ((t != NULL) ? t->mCount : 0);
         // (b (or (gethash word bad) 0)))
         t = mBadTokens.get(word);
         double b = ((t != NULL) ? t->mCount : 0);
         if ((g + b) > 5) {
             // (max .01
             //      (min .99 (float (/ (min 1 (/ b nbad))
             //                         (+ (min 1 (/ g ngood))
             //                            (min 1 (/ b nbad)))))))
             token.mProbability = dmax(.01,
                                      dmin(.99,
                                          (dmin(1.0, (b / nbad)) /
                                               (dmin(1.0, (g / ngood)) +
                                                dmin(1.0, (b / nbad))))));
         } else {
             token.mProbability = 0.4;
         }
     }
     
     // sort the array by the distance of the token probabilities from a 50-50 value of 0.5.
     PRUint32 first, last = count;
     if (count > 15) {
         first = count - 15;
         NS_QuickSort(tokens, count, sizeof(Token), compareTokens, NULL);
     } else {
         first = 0;
     }
 
     double prod1 = 1.0, prod2 = 1.0;
     for (i = first; i < last; ++i) {
         double value = tokens[i].mProbability;
         prod1 *= value;
         prod2 *= (1.0 - value);
     }
     double prob = (prod1 / (prod1 + prod2));
     PRBool isJunk = (prob >= 0.90);
 
     delete[] tokens;
 
     if (listener)
         listener->OnMessageClassified(messageURI, isJunk ? nsMsgJunkStatus(nsIJunkMailPlugin::JUNK) : nsMsgJunkStatus(nsIJunkMailPlugin::GOOD));
 }

 static PRBool readTokens(FILE* stream, Tokenizer& tokenizer)
 {
     PRUint32 tokenCount;
     if (readUInt32(stream, &tokenCount) != 1)
         return PR_FALSE;
 
     PRUint32 bufferSize = 4096;
     char* buffer = new char[bufferSize];
     if (!buffer) return PR_FALSE;
 
     for (PRUint32 i = 0; i < tokenCount; ++i) {
         PRUint32 count;
         if (readUInt32(stream, &count) != 1)
             break;
         PRUint32 size;
         if (readUInt32(stream, &size) != 1)
             break;
         if (size >= bufferSize) {
             delete[] buffer;
             PRUint32 newBufferSize = 2 * bufferSize;
             while (size >= newBufferSize)
                 newBufferSize *= 2;
             buffer = new char[newBufferSize];
             if (!buffer) return PR_FALSE;
             bufferSize = newBufferSize;
         }
         if (fread(buffer, size, 1, stream) != 1)
             break;
         buffer[size] = '\0';
         tokenizer.add(buffer, count);
     }
     
     delete[] buffer;
     
     return PR_TRUE;
 }

 void nsBayesianFilter::readTrainingData()
 {
     nsCOMPtr<nsILocalFile> file;
     nsresult rv = getTrainingFile(file);
     if (NS_FAILED(rv)) return;
     
     PRBool exists;
     rv = file->Exists(&exists);
     if (NS_FAILED(rv) || !exists) return;
 
     // open the file, and write out training data using fprintf for now.
     FILE* stream;
     rv = file->OpenANSIFileDesc("rb", &stream);
     if (NS_FAILED(rv)) return;
 
     // FIXME:  should make sure that the tokenizers are empty.
     char cookie[4];
     if (!((fread(cookie, sizeof(cookie), 1, stream) == 1) &&
           (memcmp(cookie, kMagicCookie, sizeof(cookie)) == 0) &&
           (readUInt32(stream, &mGoodCount) == 1) &&
           (readUInt32(stream, &mBadCount) == 1) &&
            readTokens(stream, mGoodTokens) &&
            readTokens(stream, mBadTokens))) {
         NS_WARNING("failed to read training data.");
     }
     
     fclose(stream);
 }