<commit_before>#include "IndriPruneIndex.h"

#include <algorithm>
#include <string>
#include <sstream>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include <boost/algorithm/string.hpp>
#include <boost/bind.hpp>
#include <boost/dynamic_bitset.hpp>
#include <boost/foreach.hpp>
#include <boost/format.hpp>
#include <boost/function.hpp>

// FIXME: check if this quick fix is appropriate
#define __atomic_add __gnu_cxx::__atomic_add

#include "indri/CompressedCollection.hpp"
#include "indri/DeletedDocumentList.hpp"
#include "indri/DiskIndex.hpp"
#include "indri/IndexWriter.hpp"
#include "indri/Parameters.hpp"
#include "indri/Path.hpp"
#include "indri/Repository.hpp"
#include "indri/TermScoreFunctionFactory.hpp"

#include "BinnedFilter.h"
#include "DocumentCentricTermScoreFunction.h"
#include "FilteredIndex.h"
#include "GainFunction.h"
#include "QuantileSampler.h"
#include "RunningStat.h"

indri_contrib::DocumentCentricTermScoreFunction*
getDocumentCentricTermScoreFunction(const std::string& stringSpec, 
				    double documentLength, double collectionLength, int documentCount)
{
    // Notation:
    //   collectionLength    := N  (collection length, in # of words)
    //   documentCount       := D  (total # of documents)
    //   documentLength      := DL (document length, in # of words)
    //
    // When invoked, pass 3 parameters:
    //   termCount           := TF (term frequency in some document)
    //   collectionTermCount := CF (collection term frequency)
    //   documentFrequency   := DF (document frequency, in # of docs)
    indri::api::Parameters spec;
    parseSpec(spec, stringSpec);
    std::string method = spec.get("method", "dirichlet");

    if (method == "dirichlet" || method == "d" || method == "dir") {
	// dirichlet -- takes parameter "mu"
	double mu = spec.get("mu", 2500);
	return new indri_contrib::DirichletDocumentCentricTermScoreFunction(
	    mu, documentLength, collectionLength);
    } 
    else if (method == "linear" || method == "jm" || method == "jelinek-mercer") {
	// jelinek-mercer -- can take parameters collectionLambda (or just lambda) and documentLambda
	double collectionLambda = spec.exists("collectionLambda")?
	    spec.get("collectionLambda", 0.4): spec.get("lambda", 0.4);
	return new indri_contrib::JelinekMercerDocumentCentricTermScoreFunction(
	    collectionLambda, documentLength, collectionLength);
    }
    else if (method == "tfidf" || method == "okapi") {
	double k1 = spec.get("k1", 1.2);
	double b = spec.get("b", 0.75);
	double k3 = spec.get("k3", 7);
	int qtf = spec.get("qtf", 1);

	return new indri_contrib::TFIDFDocumentCentricTermScoreFunction(
	    qtf, k1, b, k3, (method == "okapi"), documentLength, collectionLength, documentCount);
    } 
    else if (method == "kld") {
	double collectionLambda = spec.exists("collectionLambda")?
	    spec.get("collectionLambda", 0.0): spec.get("lambda", 0.0);
	return new indri_contrib::KLDivergenceDocumentCentricTermScoreFunction(
	    collectionLambda, documentLength, collectionLength);
    }
    else die(boost::format("No support for method '%s'") % method);
}

// Weight
//
typedef std::vector<double> Weight;

// LazyVectorLoader
//
class LazyVectorLoader {
    std::ifstream _in;
    std::string _buf;
    std::istringstream _linein;

public:
    LazyVectorLoader(std::string guardPath): _in(guardPath) { }

    bool hasNext() const { return _in; }

    void next(std::unordered_set<int>& v) {
	v.clear();

	if (getline(_in, _buf)) {
	    _linein.str(_buf);
	    _linein.clear();

	    int n;
	    while (_linein >> n) v.insert(n);
	}
    }
};


// Generalized uniform pruning<commit_msg>Fixed the problem with implicit casting - g++ not happy with that<commit_after>#include "IndriPruneIndex.h"

#include <algorithm>
#include <string>
#include <sstream>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include <boost/algorithm/string.hpp>
#include <boost/bind.hpp>
#include <boost/dynamic_bitset.hpp>
#include <boost/foreach.hpp>
#include <boost/format.hpp>
#include <boost/function.hpp>

// FIXME: check if this quick fix is appropriate
#define __atomic_add __gnu_cxx::__atomic_add

#include "indri/CompressedCollection.hpp"
#include "indri/DeletedDocumentList.hpp"
#include "indri/DiskIndex.hpp"
#include "indri/IndexWriter.hpp"
#include "indri/Parameters.hpp"
#include "indri/Path.hpp"
#include "indri/Repository.hpp"
#include "indri/TermScoreFunctionFactory.hpp"

#include "BinnedFilter.h"
#include "DocumentCentricTermScoreFunction.h"
#include "FilteredIndex.h"
#include "GainFunction.h"
#include "QuantileSampler.h"
#include "RunningStat.h"

indri_contrib::DocumentCentricTermScoreFunction*
getDocumentCentricTermScoreFunction(const std::string& stringSpec, 
				    double documentLength, double collectionLength, int documentCount)
{
    // Notation:
    //   collectionLength    := N  (collection length, in # of words)
    //   documentCount       := D  (total # of documents)
    //   documentLength      := DL (document length, in # of words)
    //
    // When invoked, pass 3 parameters:
    //   termCount           := TF (term frequency in some document)
    //   collectionTermCount := CF (collection term frequency)
    //   documentFrequency   := DF (document frequency, in # of docs)
    indri::api::Parameters spec;
    parseSpec(spec, stringSpec);
    std::string method = spec.get("method", "dirichlet");

    if (method == "dirichlet" || method == "d" || method == "dir") {
	// dirichlet -- takes parameter "mu"
	double mu = spec.get("mu", 2500);
	return new indri_contrib::DirichletDocumentCentricTermScoreFunction(
	    mu, documentLength, collectionLength);
    } 
    else if (method == "linear" || method == "jm" || method == "jelinek-mercer") {
	// jelinek-mercer -- can take parameters collectionLambda (or just lambda) and documentLambda
	double collectionLambda = spec.exists("collectionLambda")?
	    spec.get("collectionLambda", 0.4): spec.get("lambda", 0.4);
	return new indri_contrib::JelinekMercerDocumentCentricTermScoreFunction(
	    collectionLambda, documentLength, collectionLength);
    }
    else if (method == "tfidf" || method == "okapi") {
	double k1 = spec.get("k1", 1.2);
	double b = spec.get("b", 0.75);
	double k3 = spec.get("k3", 7);
	int qtf = spec.get("qtf", 1);

	return new indri_contrib::TFIDFDocumentCentricTermScoreFunction(
	    qtf, k1, b, k3, (method == "okapi"), documentLength, collectionLength, documentCount);
    } 
    else if (method == "kld") {
	double collectionLambda = spec.exists("collectionLambda")?
	    spec.get("collectionLambda", 0.0): spec.get("lambda", 0.0);
	return new indri_contrib::KLDivergenceDocumentCentricTermScoreFunction(
	    collectionLambda, documentLength, collectionLength);
    }
    else die(boost::format("No support for method '%s'") % method);
}

// Weight
//
typedef std::vector<double> Weight;

// LazyVectorLoader
//
class LazyVectorLoader {
    std::ifstream _in;
    std::string _buf;
    std::istringstream _linein;

public:
    LazyVectorLoader(std::string guardPath): _in(guardPath) { }

    bool hasNext() const { return static_cast<bool>(_in); }

    void next(std::unordered_set<int>& v) {
	v.clear();

	if (getline(_in, _buf)) {
	    _linein.str(_buf);
	    _linein.clear();

	    int n;
	    while (_linein >> n) v.insert(n);
	}
    }
};


// Generalized uniform pruning
{'fraction_non_alphanumeric': 0.13855339937576333, 'fraction_numerical': 0.007056588410910571, 'mean_word_length': 4.260528194147038, 'pattern_counts': {'":': 0, '<': 32, '<?xml version=': 0, '>': 36, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2096894', 'n_tokens_mistral': 2417, 'n_tokens_neox': 2326, 'n_words': 563}