bm25scorer.java

来自「MG4J (Managing Gigabytes for Java) is a 」· Java 代码 · 共 201 行

package it.unimi.dsi.mg4j.search.score;

/*		 
 * MG4J: Managing Gigabytes for Java
 *
 * Copyright (C) 2006-2007 Sebastiano Vigna
 *
 *  This library is free software; you can redistribute it and/or modify it
 *  under the terms of the GNU Lesser General Public License as published by the Free
 *  Software Foundation; either version 2.1 of the License, or (at your option)
 *  any later version.
 *
 *  This library is distributed in the hope that it will be useful, but
 *  WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
 *  or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
 *  for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 *
 */

import it.unimi.dsi.fastutil.ints.IntList;
import it.unimi.dsi.mg4j.index.Index;
import it.unimi.dsi.mg4j.search.DocumentIterator;
import it.unimi.dsi.mg4j.search.visitor.CounterCollectionVisitor;
import it.unimi.dsi.mg4j.search.visitor.CounterSetupVisitor;
import it.unimi.dsi.mg4j.search.visitor.TermCollectionVisitor;

import java.io.IOException;
import java.util.Arrays;

import org.apache.log4j.Logger;

/** A scorer that implements the BM25 ranking formula.
 * 
 * <p><strong>Warning</strong>: the default values {@link #DEFAULT_K1} and {@link #DEFAULT_B}
 * have changed in MG4J 1.1.2 (see below).
 * 
 * <p>BM25 is the name of a formula derived from the probabilistic model. The essential feature
 * of the formula is that of assigning to each term appearing in a given document a weight depending
 * both on the count (the number of occurrences of the term in the document), on the frequency (the
 * number of the documents in which the term appears) and on the document length (in words). 
 * 
 * <p>There are a number
 * of incarnations with small variations of the formula itself. Here, the weight
 * assigned to a term which appears in <var>f</var> documents out of a collection of <var>N</var> documents
 * w.r.t. to a document of length <var>l</var> in which the term appears <var>c</var> times is
 * <div style="text-align: center">
 * log<big>(</big> (<var>N</var> &minus; <var>f</var> + 1/2) / (f + 1/2) <big>)</big> ( <var>k</var><sub>1</sub> + 1 ) <var>c</var> &nbsp;<big>&frasl;</big>&nbsp; <big>(</big> <var>c</var> + <var>k</var><sub>1</sub> ((1 &minus; <var>b</var>) + <var>b</var><var>l</var> / <var>L</var>) <big>)</big>,
 * </div>
 * where <var>L</var> is the average document length, and <var>k</var><sub>1</sub> and <var>b</var> are
 * parameters that default to {@link #DEFAULT_K1} and {@link #DEFAULT_B}: these values were chosen
 * following the suggestions given in 
 * &ldquo;Efficiency vs. effectiveness in Terabyte-scale information retrieval&rdquo;, by Stefan B&#252;ttcher and Charles L. A. Clarke, 
 * in <i>Proceedings of the 14th Text REtrieval 
 * Conference (TREC 2005)</i>. Gaithersburg, USA, November 2005. The logarithmic part (a.k.a.
 * <em>idf (inverse document-frequency)</em> part) is actually
 * maximised with {@link #EPSILON_SCORE}, so it is never negative (the net effect being that terms appearing 
 * in more than half of the documents have almost no weight).
 *  
 * <p>This class uses a {@link it.unimi.dsi.mg4j.search.visitor.CounterCollectionVisitor}
 * and related classes to take into consideration only terms that are actually involved
 * in the current document. 
 * 
 * @author Mauro Mereu
 * @author Sebastiano Vigna
 */
public class BM25Scorer extends AbstractWeightedScorer implements DelegatingScorer {
	private static final Logger LOGGER = Logger.getLogger( BM25Scorer.class );
	private static final boolean DEBUG = false;

	/** The default value used for the parameter <var>k</var><sub>1</sub>. */
	public final static double DEFAULT_K1 = 1.2;
	/** The default value used for the parameter <var>b</var>. */
	public final static double DEFAULT_B = 0.5;
	/** The value of the document-frequency part for terms appearing in more than half of the documents. */
	public final static double EPSILON_SCORE = 1.000000082240371E-9; // TODO: put a better value
	
	/** The counter collection visitor used to estimate counts. */
	private final CounterCollectionVisitor counterCollectionVisitor;
	/** The counter setup visitor used to estimate counts. */
	private final CounterSetupVisitor setupVisitor;
	/** The term collection visitor used to estimate counts. */
	private final TermCollectionVisitor termVisitor;

	/** The parameter <var>k</var><sub>1</sub>. */
	public final double k1;
	/** The parameter <var>b</var>. */
	public final double b;

	/** The parameter {@link #k1} plus one, precomputed. */
	private final double k1Plus1;
	/** One minus {@link #b}, precomputed. */
	private final double oneMinusB;

	/** An array (parallel to {@link #currIndex}) that caches average document sizes. */
	private double averageDocumentSize[];
	/** An array (parallel to {@link #currIndex}) that caches size lists. */
	private IntList sizes[];
	/** An array (parallel to {@link #currIndex}) used by {@link #score()} to cache the current document sizes. */
	private int[] size;
	/** An array indexed by offsets that caches the inverse document-frequency part of the formula, multiplied by the index weight. */
	private double[] weightedIdfPart;
	
	public BM25Scorer() {
		this( DEFAULT_K1, DEFAULT_B );
	}

	public BM25Scorer( final double k1, final double b ) {
		termVisitor = new TermCollectionVisitor();
		setupVisitor = new CounterSetupVisitor( termVisitor );
		counterCollectionVisitor = new CounterCollectionVisitor( setupVisitor );
		this.k1 = k1;
		this.b = b;
		k1Plus1 = k1 + 1;
		oneMinusB = 1 - b;
	}

	public BM25Scorer( final String k1, final String b ) {
		this( Double.parseDouble( k1 ), Double.parseDouble( b ) );
	}

	public synchronized BM25Scorer copy() {
		final BM25Scorer scorer = new BM25Scorer( k1, b );
		scorer.setWeights( index2Weight );
		return scorer;
	}

	public double score() throws IOException {
		setupVisitor.clear();
		documentIterator.acceptOnTruePaths( counterCollectionVisitor );
		
		final int document = documentIterator.document();
		final int[] count = setupVisitor.count;
		final int[] indexNumber = setupVisitor.indexNumber;
		final double[] weightedIdfPart = this.weightedIdfPart;
		final double[] averageDocumentSize = this.averageDocumentSize;
		final int[] size = this.size;
		
		for( int i = currIndex.length; i-- != 0; ) size[ i ] = sizes[ i ].getInt( document );

		int k;
		double score = 0;
		for ( int i = count.length; i-- != 0; ) {
			k = indexNumber[ i ];
			score += ( k1Plus1 * count[ i ] ) / ( count[ i ] + k1 * ( oneMinusB + b * size[ k ] / averageDocumentSize[ k ] ) ) * weightedIdfPart[ i ];
		}
		return score;
	}

	public double score( final Index index ) {
		throw new UnsupportedOperationException();
	}


	public void wrap( DocumentIterator d ) throws IOException {
		documentIterator = d;
		termVisitor.prepare();
		d.accept( termVisitor );

		if ( DEBUG ) LOGGER.debug( "Term Visitor found " + termVisitor.numberOfPairs() + " leaves" );

		// Note that we use the index array provided by the visitor, *not* by the iterator.
		final Index[] index = termVisitor.indices();

		if ( DEBUG ) LOGGER.debug( "Indices: " + Arrays.toString( index ) );

		// Some caching of frequently-used values
		sizes = new IntList[ index.length ];
		for( int i = index.length; i-- != 0; )
			if ( ( sizes[ i ] = index[ i ].sizes ) == null ) throw new IllegalStateException( "A BM25 scorer requires document sizes" );
		
		averageDocumentSize = new double[ index.length ];
		for ( int i = index.length; i-- != 0; )
			averageDocumentSize[ i ] = (double)( index[ i ].numberOfOccurrences ) / index[ i ].numberOfDocuments;

		if ( DEBUG ) LOGGER.debug( "Average document sizes: " + Arrays.toString( averageDocumentSize ) );
		setupVisitor.prepare();
		d.accept( setupVisitor );

		final int[] frequency = setupVisitor.frequency;
		final int[] indexNumber = setupVisitor.indexNumber;
		
		// We do all logs here, and multiply by the weight
		weightedIdfPart = new double[ frequency.length ];
		for( int i = weightedIdfPart.length; i-- != 0; )
			weightedIdfPart[ i ] = Math.max( EPSILON_SCORE,  
					Math.log( ( index[ indexNumber[ i ] ].numberOfDocuments - frequency[ i ] + 0.5 ) / ( frequency[ i ] + 0.5 ) ) ) * index2Weight.getDouble( index[ indexNumber[ i ] ] );
			
		size = new int[ index.length ];
		currIndex = index;
	}
	
	public boolean usesIntervals() {
		return false;
	}

}