BagOfWords.py

from __future__ import division
import nltk
import string
from TFIDF import tfidf

# "Adapting a technique of Das and Chen (2001), we added the tag NOT to every word between a negation word ('not',
# 'isn't', 'didn't', etc.) and the first punctuation mark following the negation word."
# They didn't provide a full list.
NEGATION_WORDS      = ["not", "isn't", "didn't", "doesn't"]
PUNCTUATION         = [".", "!", "?", ",", ";"] #TODO make this work with POS tags (._.)
POSITION_TAGS       = ["_1Q", "_2H", "_3Q"]
POSITION_THRESHOLDS = [0.25, 0.75, 1]

# ref_bag is used to calculate the total word count across all documents.
def make(text, ref_bag=None, gram_length=1, use_negation=False, use_presence=False, use_pos_tags=False, use_adj_only=False, use_position=False, normalize=True, use_hash=True):
  bag_of_words = {}
  if use_negation:
    do_negation = False
  
  words = nltk.word_tokenize(text)
  if use_pos_tags:
    tagged = nltk.pos_tag(words)
    words = [string.join(t, "_") for t in tagged]
  for i in range(len(words) - gram_length + 1):
    n_gram = string.join(words[i:i+gram_length], "_")
    
    if use_negation:
      if (gram_length == 1): # Pang and Lee didn't do negation tagging for bigrams.
        if n_gram in NEGATION_WORDS:
          do_negation = True
        elif n_gram in PUNCTUATION:
          do_negation = False
        if do_negation:
          n_gram = "NOT_" + n_gram
          
    if use_position:
      for j in range(len(POSITION_TAGS)):
        if i/len(words) < POSITION_THRESHOLDS[j]:
          n_gram += POSITION_TAGS[j]
          break
          
    # LIBSVM won't use strings as keys, so hash to convert to a number.
    if use_hash:
      index = hash(n_gram)
    else:
      index = n_gram
      
    if not (use_pos_tags and use_adj_only and (tagged[i][1] != "JJ")):
      if (not use_presence) and bag_of_words.has_key(index):
        bag_of_words[index] += 1
      else:
        bag_of_words[index] = 1
      
      # Add it to the reference bag
      if ref_bag != None:
        if ref_bag.has_key(index):
          ref_bag[index] += 1
        else:
          ref_bag[index] = 1
          
  if normalize:
    length = 0
    for k in bag_of_words.keys():
      length += (bag_of_words[k]**2)
    length **= 0.5
    for k in bag_of_words.keys():
      bag_of_words[k] = bag_of_words[k]/length
  return bag_of_words

# document and document are lists of words (pre-tokenized with nltk.word_tokenize())
def make_tfidf(document, documents):
  bag = {}
  factor = 0
  for term in set(document):
    weight = tfidf(term, document, documents)
    if (weight != 0):
      bag[term] = weight
      factor += weight**2
  factor **= 0.5
  for key in bag.keys():
    bag[key] /= factor
  return bag