<NME> tfidfmodel.py
<BEF> #!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright (C) 2012 Radim Rehurek <radimrehurek@seznam.cz>
# Copyright (C) 2017 Mohit Rathore <mrmohitrathoremr@gmail.com>
# Licensed under the GNU LGPL v2.1 - http://www.gnu.org/licenses/lgpl.html

"""This module implements functionality related to the `Term Frequency - Inverse Document Frequency
<https://en.wikipedia.org/wiki/Tf%E2%80%93idf>`_ class of bag-of-words vector space models.

"""

import logging
from functools import partial
import re
logger = logging.getLogger('gensim.models.tfidfmodel')


class TfidfModel(interfaces.TransformationABC):
    """
    Objects of this class realize the transformation between word-document co-occurence
    matrix (integers) into a locally/globally weighted matrix (positive floats).

    This is done by combining the term frequency counts (the TF part) with inverse
    document frequency counts (the IDF part), optionally normalizing the resulting
    documents to unit length.

    The main methods are:

    1. constructor, which calculates IDF weights for all terms in the training corpus.
    2. the [] method, which transforms a simple count representation into the TfIdf
       space.

            * `a` - augmented,
            * `l` - logarithm,
            * `d` - double logarithm,

    Model persistency is achieved via its load/save methods.
    """
    def __init__(self, corpus=None, id2word=None, dictionary=None, normalize_tfidf=True, normalize_tf=False, log_base=2, add=0):
        """
        `normalize` dictates whether the transformed vectors will be set to unit
        length.

            "You must train two separate tf-idf models.".format(
        and it will be used to directly construct the inverse document frequency
        mapping (then `corpus`, if specified, is ignored).
        """
        self.normalize_tf = normalize_tf
        self.normalize_tfidf = normalize_tfidf
        self.log_base = log_base
        self.add = add
        self.id2word = id2word
        self.num_docs, self.num_nnz, self.idfs = None, None, None
        if dictionary is not None:
            if corpus is not None:
                logger.warning("constructor received both corpus and explicit "
                               "inverse document frequencies; ignoring the corpus")
            self.num_docs, self.num_nnz = dictionary.num_docs, dictionary.num_nnz
            self.idfs = self.dfs2idfs(dictionary.dfs, dictionary.num_docs)
        elif corpus is not None:
            self.initialize(corpus)
        else:

    # resolve aliases
            pass


    def dfs2idfs(self, dfs, totaldocs):
        """
        Given a mapping of `term->document frequency`, construct a mapping of
        `term->inverse document frequency`.
        """
        return dict((termid, self.add + math.log(1.0 * totaldocs / docfreq, self.log_base))
                    for termid, docfreq in dfs.iteritems())

    def idfs2dfs(self, idfs, totaldocs):
        """
        Inverse mapping for `dfs2idfs`.
        """
        return dict((termid, int(round(totaldocs / self.log_base**(weight-self.add)))
                    for termid, weight in idfs.iteritems())

    def __str__(self):
        return "TfidfModel(num_docs=%s, num_nnz=%s)" % (self.num_docs, self.num_nnz)


    return w_tf + w_df + w_n


def df2idf(docfreq, totaldocs, log_base=2.0, add=0.0):
    r"""Compute inverse-document-frequency for a term with the given document frequency `docfreq`:
    :math:`idf = add + log_{log\_base} \frac{totaldocs}{docfreq}`

    Parameters
    ----------
            if docno % 10000 == 0:
                logger.info("PROGRESS: processing document #%i" % docno)
            numnnz += len(bow)
            for termid, termcount in bow:
                dfs[termid] = dfs.get(termid, 0) + 1

        # keep some stats about the training corpus
        self.num_docs = docno + 1 # HACK using leftover from enumerate(corpus) above
        self.num_nnz = numnnz

        # and finally compute the idf weights
        logger.info("calculating IDF weights for %i documents and %i features (%i matrix non-zeros)" %
                     (self.num_docs, 1 + max([-1] + dfs.keys()), self.num_nnz))
        self.idfs = self.dfs2idfs(dfs, self.num_docs)


    def __getitem__(self, bow):
def precompute_idfs(wglobal, dfs, total_docs):
    """Pre-compute the inverse document frequency mapping for all terms.

    Parameters
    ----------
        if is_corpus:
            return self._apply(bow)

        if self.normalize_tf == 'doclen':
            doc_length = sum(tf for _, tf in bow)
            bow = [(termid, float(tf)/doc_length) for termid, tf in bow]
        if self.normalize_tf == 'sqrt':
            bow = [(termid, math.sqrt(tf)) for termid, tf in bow]


        # unknown (new) terms will be given zero weight (NOT infinity/huge weight,
        # as strict application of the IDF formula would dictate)
        vector = [(termid, tf * self.idfs.get(termid, 0.0))
                  for termid, tf in bow if self.idfs.get(termid, 0.0) != 0.0]
        if self.normalize_tfidf == 'unit':
            vector = matutils.unitvec(vector)
        return vector
#endclass TfidfModel
        Inverse document frequencies in the format `{term_id_1: idfs_1, term_id_2: idfs_2, ...}`.

    """
    # not strictly necessary and could be computed on the fly in TfidfModel__getitem__.
    # this method is here just to speed things up a little.
    return {termid: wglobal(df, total_docs) for termid, df in dfs.items()}


def smartirs_wlocal(tf, local_scheme):
    """Calculate local term weight for a term using the weighting scheme specified in `local_scheme`.

    Parameters
    ----------
    tf : int
        Term frequency.
    local : {'b', 'n', 'a', 'l', 'd', 'L'}
        Local transformation scheme.

    Returns
    -------
    float
        Calculated local weight.

    """
    if local_scheme == "n":
        return tf
    elif local_scheme == "l":
        return 1 + np.log2(tf)
    elif local_scheme == "d":
        return 1 + np.log2(1 + np.log2(tf))
    elif local_scheme == "a":
        return 0.5 + (0.5 * tf / tf.max(axis=0))
    elif local_scheme == "b":
        return tf.astype('bool').astype('int')
    elif local_scheme == "L":
        return (1 + np.log2(tf)) / (1 + np.log2(tf.mean(axis=0)))


def smartirs_wglobal(docfreq, totaldocs, global_scheme):
    """Calculate global document weight based on the weighting scheme specified in `global_scheme`.

    Parameters
    ----------
    docfreq : int
        Document frequency.
    totaldocs : int
        Total number of documents.
    global_scheme : {'n', 'f', 't', 'p'}
        Global transformation scheme.

    Returns
    -------
    float
        Calculated global weight.

    """
    if global_scheme == "n":
        return 1.0
    elif global_scheme == "f":
        return np.log2(1.0 * totaldocs / docfreq)
    elif global_scheme == "t":
        return np.log2((totaldocs + 1.0) / docfreq)
    elif global_scheme == "p":
        return max(0, np.log2((1.0 * totaldocs - docfreq) / docfreq))


@deprecated("Function will be removed in 4.0.0")
def smartirs_normalize(x, norm_scheme, return_norm=False):
    """Normalize a vector using the normalization scheme specified in `norm_scheme`.

    Parameters
    ----------
    x : numpy.ndarray
        The tf-idf vector.
    norm_scheme : {'n', 'c'}
        Document length normalization scheme.
    return_norm : bool, optional
        Return the length of `x` as well?

    Returns
    -------
    numpy.ndarray
        Normalized array.
    float (only if return_norm is set)
        Norm of `x`.
    """
    if norm_scheme == "n":
        if return_norm:
            _, length = matutils.unitvec(x, return_norm=return_norm)
            return x, length
        else:
            return x
    elif norm_scheme == "c":
        return matutils.unitvec(x, return_norm=return_norm)


class TfidfModel(interfaces.TransformationABC):
    """Objects of this class realize the transformation between word-document co-occurrence matrix (int)
    into a locally/globally weighted TF-IDF matrix (positive floats).

    Examples
    --------
    .. sourcecode:: pycon

        >>> import gensim.downloader as api
        >>> from gensim.models import TfidfModel
        >>> from gensim.corpora import Dictionary
        >>>
        >>> dataset = api.load("text8")
        >>> dct = Dictionary(dataset)  # fit dictionary
        >>> corpus = [dct.doc2bow(line) for line in dataset]  # convert corpus to BoW format
        >>>
        >>> model = TfidfModel(corpus)  # fit model
        >>> vector = model[corpus[0]]  # apply model to the first corpus document

    """
    def __init__(self, corpus=None, id2word=None, dictionary=None, wlocal=utils.identity,
                 wglobal=df2idf, normalize=True, smartirs=None, pivot=None, slope=0.25):
        r"""Compute TF-IDF by multiplying a local component (term frequency) with a global component
        (inverse document frequency), and normalizing the resulting documents to unit length.
        Formula for non-normalized weight of term :math:`i` in document :math:`j` in a corpus of :math:`D` documents

        .. math:: weight_{i,j} = frequency_{i,j} * log_2 \frac{D}{document\_freq_{i}}

        or, more generally

        .. math:: weight_{i,j} = wlocal(frequency_{i,j}) * wglobal(document\_freq_{i}, D)

        so you can plug in your own custom :math:`wlocal` and :math:`wglobal` functions.

        Parameters
        ----------
        corpus : iterable of iterable of (int, int), optional
            Input corpus
        id2word : {dict, :class:`~gensim.corpora.Dictionary`}, optional
            Mapping token - id, that was used for converting input data to bag of words format.
        dictionary : :class:`~gensim.corpora.Dictionary`
            If `dictionary` is specified, it must be a `corpora.Dictionary` object and it will be used.
            to directly construct the inverse document frequency mapping (then `corpus`, if specified, is ignored).
        wlocals : callable, optional
            Function for local weighting, default for `wlocal` is :func:`~gensim.utils.identity`
            (other options: :func:`numpy.sqrt`, `lambda tf: 0.5 + (0.5 * tf / tf.max())`, etc.).
        wglobal : callable, optional
            Function for global weighting, default is :func:`~gensim.models.tfidfmodel.df2idf`.
        normalize : {bool, callable}, optional
            Normalize document vectors to unit euclidean length? You can also inject your own function into `normalize`.
        smartirs : str, optional
            SMART (System for the Mechanical Analysis and Retrieval of Text) Information Retrieval System,
            a mnemonic scheme for denoting tf-idf weighting variants in the vector space model.
            The mnemonic for representing a combination of weights takes the form XYZ,
            for example 'ntc', 'bpn' and so on, where the letters represents the term weighting of the document vector.

            Term frequency weighing:
                * `b` - binary,
                * `t` or `n` - raw,
                * `a` - augmented,
                * `l` - logarithm,
                * `d` - double logarithm,
                * `L` - log average.

            Document frequency weighting:
                * `x` or `n` - none,
                * `f` - idf,
                * `t` - zero-corrected idf,
                * `p` - probabilistic idf.

            Document normalization:
                * `x` or `n` - none,
                * `c` - cosine,
                * `u` - pivoted unique,
                * `b` - pivoted character length.

            Default is 'nfc'.
            For more information visit `SMART Information Retrieval System
            <https://en.wikipedia.org/wiki/SMART_Information_Retrieval_System>`_.
        pivot : float or None, optional
            In information retrieval, TF-IDF is biased against long documents [1]_. Pivoted document length
            normalization solves this problem by changing the norm of a document to `slope * old_norm + (1.0 -
            slope) * pivot`.

            You can either set the `pivot` by hand, or you can let Gensim figure it out automatically with the following
            two steps:

                * Set either the `u` or `b` document normalization in the `smartirs` parameter.
                * Set either the `corpus` or `dictionary` parameter. The `pivot` will be automatically determined from
                  the properties of the `corpus` or `dictionary`.

            If `pivot` is None and you don't follow steps 1 and 2, then pivoted document length normalization will be
            disabled. Default is None.

            See also the blog post at https://rare-technologies.com/pivoted-document-length-normalisation/.
        slope : float, optional
            In information retrieval, TF-IDF is biased against long documents [1]_. Pivoted document length
            normalization solves this problem by changing the norm of a document to `slope * old_norm + (1.0 -
            slope) * pivot`.

            Setting the `slope` to 0.0 uses only the `pivot` as the norm, and setting the `slope` to 1.0 effectively
            disables pivoted document length normalization. Singhal [2]_ suggests setting the `slope` between 0.2 and
            0.3 for best results. Default is 0.25.

            See also the blog post at https://rare-technologies.com/pivoted-document-length-normalisation/.

        References
        ----------
        .. [1] Singhal, A., Buckley, C., & Mitra, M. (1996). `Pivoted Document Length
           Normalization <http://singhal.info/pivoted-dln.pdf>`_. *SIGIR Forum*, 51, 176–184.
        .. [2] Singhal, A. (2001). `Modern information retrieval: A brief overview <http://singhal.info/ieee2001.pdf>`_.
           *IEEE Data Eng. Bull.*, 24(4), 35–43.

        """
        self.id2word = id2word
        self.wlocal, self.wglobal, self.normalize = wlocal, wglobal, normalize
        self.num_docs, self.num_nnz, self.idfs = None, None, None
        self.smartirs = resolve_weights(smartirs) if smartirs is not None else None
        self.slope = slope
        self.pivot = pivot
        self.eps = 1e-12

        if smartirs:
            n_tf, n_df, n_n = self.smartirs
            self.wlocal = partial(smartirs_wlocal, local_scheme=n_tf)
            self.wglobal = partial(smartirs_wglobal, global_scheme=n_df)

        if dictionary:
            # user supplied a Dictionary object, which already contains all the
            # statistics we need to construct the IDF mapping. we can skip the
            # step that goes through the corpus (= an optimization).
            if corpus:
                logger.warning(
                    "constructor received both corpus and explicit inverse document frequencies; ignoring the corpus"
                )
            self.num_docs, self.num_nnz = dictionary.num_docs, dictionary.num_nnz
            self.cfs = dictionary.cfs.copy()
            self.dfs = dictionary.dfs.copy()
            self.term_lens = {termid: len(term) for termid, term in dictionary.items()}
            self.idfs = precompute_idfs(self.wglobal, self.dfs, self.num_docs)
            if not id2word:
                self.id2word = dictionary
        elif corpus:
            self.initialize(corpus)
        else:
            # NOTE: everything is left uninitialized; presumably the model will
            # be initialized in some other way
            pass

        # If smartirs is not None, override pivot and normalize
        if not smartirs:
            return
        if self.pivot is not None:
            if n_n in 'ub':
                logger.warning("constructor received pivot; ignoring smartirs[2]")
            return
        if n_n in 'ub' and callable(self.normalize):
            logger.warning("constructor received smartirs; ignoring normalize")
        if n_n in 'ub' and not dictionary and not corpus:
            logger.warning("constructor received no corpus or dictionary; ignoring smartirs[2]")
        elif n_n == "u":
            self.pivot = 1.0 * self.num_nnz / self.num_docs
        elif n_n == "b":
            self.pivot = 1.0 * sum(
                self.cfs[termid] * (self.term_lens[termid] + 1.0) for termid in dictionary.keys()
            ) / self.num_docs

    @classmethod
    def load(cls, *args, **kwargs):
        """Load a previously saved TfidfModel class. Handles backwards compatibility from
        older TfidfModel versions which did not use pivoted document normalization.

        """
        model = super(TfidfModel, cls).load(*args, **kwargs)
        if not hasattr(model, 'pivot'):
            model.pivot = None
            logger.info('older version of %s loaded without pivot arg', cls.__name__)
            logger.info('Setting pivot to %s.', model.pivot)
        if not hasattr(model, 'slope'):
            model.slope = 0.65
            logger.info('older version of %s loaded without slope arg', cls.__name__)
            logger.info('Setting slope to %s.', model.slope)
        if not hasattr(model, 'smartirs'):
            model.smartirs = None
            logger.info('older version of %s loaded without smartirs arg', cls.__name__)
            logger.info('Setting smartirs to %s.', model.smartirs)
        return model

    def __str__(self):
        return "%s<num_docs=%s, num_nnz=%s>" % (self.__class__.__name__, self.num_docs, self.num_nnz)

    def initialize(self, corpus):
        """Compute inverse document weights, which will be used to modify term frequencies for documents.

        Parameters
        ----------
        corpus : iterable of iterable of (int, int)
            Input corpus.

        """
        logger.info("collecting document frequencies")
        dfs = {}
        numnnz, docno = 0, -1

        for docno, bow in enumerate(corpus):
            if docno % 10000 == 0:
                logger.info("PROGRESS: processing document #%i", docno)
            numnnz += len(bow)
            for termid, _ in bow:
                dfs[termid] = dfs.get(termid, 0) + 1
        # keep some stats about the training corpus
        self.num_docs = docno + 1
        self.num_nnz = numnnz
        self.cfs = None
        self.dfs = dfs
        self.term_lengths = None
        # and finally compute the idf weights
        self.idfs = precompute_idfs(self.wglobal, self.dfs, self.num_docs)
        self.add_lifecycle_event(
            "initialize",
            msg=(
                f"calculated IDF weights for {self.num_docs} documents and {max(dfs.keys()) + 1 if dfs else 0}"
                f" features ({self.num_nnz} matrix non-zeros)"
            ),
        )

    def __getitem__(self, bow, eps=1e-12):
        """Get the tf-idf representation of an input vector and/or corpus.

        bow : {list of (int, int), iterable of iterable of (int, int)}
            Input document in the `sparse Gensim bag-of-words format
            <https://radimrehurek.com/gensim/intro.html#core-concepts>`_,
            or a streamed corpus of such documents.
        eps : float
            Threshold value, will remove all position that have tfidf-value less than `eps`.

        Returns
        -------
        vector : list of (int, float)
            TfIdf vector, if `bow` is a single document
        :class:`~gensim.interfaces.TransformedCorpus`
            TfIdf corpus, if `bow` is a corpus.

        """
        self.eps = eps
        # if the input vector is in fact a corpus, return a transformed corpus as a result
        is_corpus, bow = utils.is_corpus(bow)
        if is_corpus:
            return self._apply(bow)

        # unknown (new) terms will be given zero weight (NOT infinity/huge weight,
        # as strict application of the IDF formula would dictate)

        termid_array, tf_array = [], []
        for termid, tf in bow:
            termid_array.append(termid)
            tf_array.append(tf)

        tf_array = self.wlocal(np.array(tf_array))

        vector = [
            (termid, tf * self.idfs.get(termid))
            for termid, tf in zip(termid_array, tf_array) if abs(self.idfs.get(termid, 0.0)) > self.eps
        ]

        # and finally, normalize the vector either to unit length, or use a
        # user-defined normalization function
        if self.smartirs:
            n_n = self.smartirs[2]
            if n_n == "n" or (n_n in 'ub' and self.pivot is None):
                if self.pivot is not None:
                    _, old_norm = matutils.unitvec(vector, return_norm=True)
                norm_vector = vector
            elif n_n == "c":
                if self.pivot is not None:
                    _, old_norm = matutils.unitvec(vector, return_norm=True)
                else:
                    norm_vector = matutils.unitvec(vector)
            elif n_n == "u":
                _, old_norm = matutils.unitvec(vector, return_norm=True, norm='unique')
            elif n_n == "b":
                old_norm = sum(freq * (self.term_lens[termid] + 1.0) for termid, freq in bow)
        else:
            if self.normalize is True:
                self.normalize = matutils.unitvec
            elif self.normalize is False:
                self.normalize = utils.identity

            if self.pivot is not None:
                _, old_norm = self.normalize(vector, return_norm=True)
            else:
                norm_vector = self.normalize(vector)

        if self.pivot is None:
            norm_vector = [(termid, weight) for termid, weight in norm_vector if abs(weight) > self.eps]
        else:
            pivoted_norm = (1 - self.slope) * self.pivot + self.slope * old_norm
            norm_vector = [
                (termid, weight / float(pivoted_norm))
                for termid, weight in vector
                if abs(weight / float(pivoted_norm)) > self.eps
            ]
        return norm_vector

<MSG> make tfidf more flexible, accepting arbitrary fncs

<DFF> @@ -16,18 +16,31 @@ from gensim import interfaces, matutils, utils
 logger = logging.getLogger('gensim.models.tfidfmodel')
 
 
-class TfidfModel(interfaces.TransformationABC):
+def df2idf(docfreq, totaldocs, log_base=2.0, add=0.0):
+    """
+    Compute default inverse-document-frequency for a term with document frequency `doc_freq`::
+
+      idf = add + log(totaldocs / doc_freq)
     """
-    Objects of this class realize the transformation between word-document co-occurence
-    matrix (integers) into a locally/globally weighted matrix (positive floats).
+    return add + math.log(1.0 * totaldocs / docfreq, log_base)
 
-    This is done by combining the term frequency counts (the TF part) with inverse
-    document frequency counts (the IDF part), optionally normalizing the resulting
-    documents to unit length.
+
+def precompute_idfs(wglobal, dfs, total_docs):
+    """Precompute the inverse document frequency mapping for all terms."""
+    # not strictly necessary and could be computed on the fly in __get__.
+    # this method is here just to speed things up a little.
+    return dict((termid, wglobal(df, total_docs))
+                for termid, df in dfs.iteritems())
+
+
+class TfidfModel(interfaces.TransformationABC):
+    """
+    Objects of this class realize the transformation between word-document co-occurrence
+    matrix (integers) into a locally/globally weighted TF_IDF matrix (positive floats).
 
     The main methods are:
 
-    1. constructor, which calculates IDF weights for all terms in the training corpus.
+    1. constructor, which calculates inverse document counts for all terms in the training corpus.
     2. the [] method, which transforms a simple count representation into the TfIdf
        space.
 
@@ -37,8 +50,25 @@ class TfidfModel(interfaces.TransformationABC):
 
     Model persistency is achieved via its load/save methods.
     """
-    def __init__(self, corpus=None, id2word=None, dictionary=None, normalize_tfidf=True, normalize_tf=False, log_base=2, add=0):
+    def __init__(self, corpus=None, id2word=None, dictionary=None,
+                 wlocal=utils.identity, wglobal=df2idf, normalize=True):
         """
+        Compute tf-idf by multiplying a local component (term frequency) with a
+        global component (inverse document frequency), and optionally normalizing
+        the resulting documents to unit length. Formula for unnormalized weight
+        of term `i` in document `j`::
+
+          weight_{i,j} = frequency_{i,j} * log_2(total_docs / document_freq_{i})
+
+        or, more generally::
+
+          weight_{i,j} = wlocal(frequency_{i,j}) * wglobal(document_freq_{i}, D)
+
+        so you can plug in your own custom `wlocal` and `wglobal` functions.
+        Default for `wlocal` is identity (other options: math.sqrt, math.log, ...)
+        and default for `wglobal` is `log_2(total_docs / doc_freq)`, giving the
+        formula above.
+
         `normalize` dictates whether the transformed vectors will be set to unit
         length.
 
@@ -46,18 +76,17 @@ class TfidfModel(interfaces.TransformationABC):
         and it will be used to directly construct the inverse document frequency
         mapping (then `corpus`, if specified, is ignored).
         """
-        self.normalize_tf = normalize_tf
-        self.normalize_tfidf = normalize_tfidf
-        self.log_base = log_base
-        self.add = add
+        self.normalize = normalize
         self.id2word = id2word
+        self.wlocal, self.wglobal = wlocal, wglobal
         self.num_docs, self.num_nnz, self.idfs = None, None, None
         if dictionary is not None:
             if corpus is not None:
                 logger.warning("constructor received both corpus and explicit "
                                "inverse document frequencies; ignoring the corpus")
             self.num_docs, self.num_nnz = dictionary.num_docs, dictionary.num_nnz
-            self.idfs = self.dfs2idfs(dictionary.dfs, dictionary.num_docs)
+            self.dfs = dictionary.dfs.copy()
+            self.idfs = precompute_idfs(self.wglobal, self.dfs, self.num_docs)
         elif corpus is not None:
             self.initialize(corpus)
         else:
@@ -66,21 +95,6 @@ class TfidfModel(interfaces.TransformationABC):
             pass
 
 
-    def dfs2idfs(self, dfs, totaldocs):
-        """
-        Given a mapping of `term->document frequency`, construct a mapping of
-        `term->inverse document frequency`.
-        """
-        return dict((termid, self.add + math.log(1.0 * totaldocs / docfreq, self.log_base))
-                    for termid, docfreq in dfs.iteritems())
-
-    def idfs2dfs(self, idfs, totaldocs):
-        """
-        Inverse mapping for `dfs2idfs`.
-        """
-        return dict((termid, int(round(totaldocs / self.log_base**(weight-self.add)))
-                    for termid, weight in idfs.iteritems())
-
     def __str__(self):
         return "TfidfModel(num_docs=%s, num_nnz=%s)" % (self.num_docs, self.num_nnz)
 
@@ -97,17 +111,18 @@ class TfidfModel(interfaces.TransformationABC):
             if docno % 10000 == 0:
                 logger.info("PROGRESS: processing document #%i" % docno)
             numnnz += len(bow)
-            for termid, termcount in bow:
+            for termid, _ in bow:
                 dfs[termid] = dfs.get(termid, 0) + 1
 
         # keep some stats about the training corpus
-        self.num_docs = docno + 1 # HACK using leftover from enumerate(corpus) above
+        self.num_docs = docno + 1
         self.num_nnz = numnnz
+        self.dfs = dfs
 
         # and finally compute the idf weights
         logger.info("calculating IDF weights for %i documents and %i features (%i matrix non-zeros)" %
                      (self.num_docs, 1 + max([-1] + dfs.keys()), self.num_nnz))
-        self.idfs = self.dfs2idfs(dfs, self.num_docs)
+        self.idfs = precompute_idfs(self.wglobal, self.dfs, self.num_docs)
 
 
     def __getitem__(self, bow):
@@ -119,18 +134,11 @@ class TfidfModel(interfaces.TransformationABC):
         if is_corpus:
             return self._apply(bow)
 
-        if self.normalize_tf == 'doclen':
-            doc_length = sum(tf for _, tf in bow)
-            bow = [(termid, float(tf)/doc_length) for termid, tf in bow]
-        if self.normalize_tf == 'sqrt':
-            bow = [(termid, math.sqrt(tf)) for termid, tf in bow]
-
-
         # unknown (new) terms will be given zero weight (NOT infinity/huge weight,
         # as strict application of the IDF formula would dictate)
-        vector = [(termid, tf * self.idfs.get(termid, 0.0))
+        vector = [(termid, self.wlocal(tf) * self.idfs.get(termid))
                   for termid, tf in bow if self.idfs.get(termid, 0.0) != 0.0]
-        if self.normalize_tfidf == 'unit':
+        if self.normalize:
             vector = matutils.unitvec(vector)
         return vector
 #endclass TfidfModel

{'addition_count': 48, 'commit_subject': 'make tfidf more flexible, accepting arbitrary fncs', 'deletion_count': 40, 'file_extension': '.py', 'lang': 'py', 'license': 'lgpl-2.1', 'repo_name': 'RaRe-Technologies/gensim'}