jaepil commented on code in PR #15948: URL: https://github.com/apache/lucene/pull/15948#discussion_r3167858182
########## lucene/core/src/java/org/apache/lucene/search/BayesianScoreEstimator.java: ########## @@ -0,0 +1,228 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.lucene.search; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; +import org.apache.lucene.index.IndexReader; +import org.apache.lucene.index.StoredFields; +import org.apache.lucene.index.Term; +import org.apache.lucene.util.ArrayUtil; + +/** + * Estimates {@link BayesianScoreQuery} parameters (alpha, beta, base rate) from corpus statistics + * via pseudo-query sampling. + * + * <p>The estimation algorithm: + * + * <ol> + * <li>Sample N documents randomly from the index + * <li>For each document, create a pseudo-query from its first few tokens in the target field + * <li>Run each pseudo-query via BM25 and collect the score distribution + * <li>Estimate: beta = median(scores), alpha = 1 / std(scores) + * <li>Estimate base rate: mean fraction of documents scoring above the 95th percentile + * </ol> + * + * @lucene.experimental + */ +public class BayesianScoreEstimator { + + /** Estimated parameters for {@link BayesianScoreQuery}. */ + public record Parameters(float alpha, float beta, float baseRate) {} + + private static final int DEFAULT_N_SAMPLES = 50; + private static final int DEFAULT_TOKENS_PER_QUERY = 5; + private static final double PERCENTILE_THRESHOLD = 0.95; + private static final float BASE_RATE_MIN = 1e-6f; + private static final float BASE_RATE_MAX = 0.5f; + + private BayesianScoreEstimator() {} + + /** + * Estimates BayesianScoreQuery parameters from the given index. + * + * @param searcher the index searcher to sample from + * @param field the text field to create pseudo-queries for + * @param nSamples number of documents to sample (default 50) + * @param tokensPerQuery number of tokens per pseudo-query (default 5) + * @param seed random seed for reproducible sampling + * @return estimated alpha, beta, and base rate + * @throws IOException if an I/O error occurs reading the index + */ + public static Parameters estimate( + IndexSearcher searcher, String field, int nSamples, int tokensPerQuery, long seed) + throws IOException { + IndexReader reader = searcher.getIndexReader(); + int maxDoc = reader.maxDoc(); + if (maxDoc == 0) { + return new Parameters(1.0f, 0.0f, 0.01f); + } + + nSamples = Math.min(nSamples, maxDoc); + Random rng = new Random(seed); + + // Sample document IDs + int[] sampledDocs = sampleDocIds(maxDoc, nSamples, rng); + + // Create pseudo-queries and collect scores + List<float[]> allScoreArrays = new ArrayList<>(); + List<Float> baseRateFractions = new ArrayList<>(); + StoredFields storedFields = reader.storedFields(); + + for (int docId : sampledDocs) { + String fieldValue = storedFields.document(docId).get(field); + if (fieldValue == null || fieldValue.isEmpty()) { + continue; + } + + // Extract first N tokens as pseudo-query terms + String[] tokens = tokenize(fieldValue, tokensPerQuery); Review Comment: You're right — taking the first N tokens biases heavily toward boilerplate prefixes (license headers, legal preambles, structured templates), and on those corpora the pseudo-queries collapse into near-duplicates. The cleaner fix, I think, is to drop the document-text path entirely and reservoir-sample over the field's indexed vocabulary via MultiTerms.getTerms(reader, field) + TermsEnum. Vocabulary-level sampling is uniform over unique terms, not over occurrences — a boilerplate term that appears in 100% of documents has the same selection probability as a rare content term, so shared-prefix corpora no longer dominate the sample. If this direction sounds right, I'll prepare a follow-up commit with a regression test for the shared-prefix case. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
