├── solr_plugins ├── src │ ├── __init__.py │ └── main │ │ ├── resources │ │ └── META-INF │ │ │ └── MANIFEST.MF │ │ └── java │ │ └── org │ │ └── dice │ │ └── solrenhancements │ │ ├── similarity │ │ ├── HammingSimilarity.java │ │ ├── PayloadOnlySimilarity.java │ │ └── DiceDefaultSimilarity.java │ │ ├── queryparsers │ │ ├── PayloadAwareExtendedDismaxQParserPlugin.java │ │ └── PayloadAwareExtendDismaxQParser.java │ │ └── JarVersion.java ├── DiceSolrEnhancements-1.0.jar └── pom.xml ├── solr_configs ├── Readme.md ├── solrconfig_configure_payloadedismax_parser.xml └── schema_xml_field_definitions.xml ├── python ├── lsh_variants.py ├── vector_thresholding.py ├── solr_encode_vectors.py ├── kmeans_tree.py ├── searchers.py └── Tutorial.ipynb ├── .gitignore ├── Readme.md └── LICENSE /solr_plugins/src/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /solr_plugins/DiceSolrEnhancements-1.0.jar: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DiceTechJobs/VectorsInSearch/HEAD/solr_plugins/DiceSolrEnhancements-1.0.jar -------------------------------------------------------------------------------- /solr_plugins/src/main/resources/META-INF/MANIFEST.MF: -------------------------------------------------------------------------------- 1 | Specification-Title: App Name 2 | Specification-Version: ${pom.version} - ${build.time} 3 | Specification-Vendor: Company Name 4 | Implementation-Title: App Name 5 | Implementation-Version: ${pom.version} - ${build.time} 6 | Implementation-Vendor: Company Name 7 | Built-By: ${user.name} 8 | Build-Jdk: ${java.version} 9 | Build-Time: ${build.time} 10 | Built-Date: ${build.time} 11 | -------------------------------------------------------------------------------- /solr_configs/Readme.md: -------------------------------------------------------------------------------- 1 | # Solr Config Snippets 2 | 3 | Due to licensing, I cannot copy the entire solr config and schema xml here. So instead I have just included the snippets needed to enable the solr plugins: 4 | 5 | ## 1. **solrconfig.xml** 6 | - Add the payloadEdismax parser to the solrconfig.xml. 7 | - Configure the solrconfig.xml to load the plugins jar file 8 | 9 | ## 2. **schema.xml** 10 | - Add the field definitions for the vector field and cluster field types 11 | - Set the similarity class to schema similarity 12 | -------------------------------------------------------------------------------- /solr_plugins/src/main/java/org/dice/solrenhancements/similarity/HammingSimilarity.java: -------------------------------------------------------------------------------- 1 | 2 | package org.dice.solrenhancements.similarity; 3 | 4 | import org.apache.lucene.search.similarities.ClassicSimilarity; 5 | 6 | /** 7 | * Created by simon.hughes on 4/16/14. 8 | * Turn off all weightings 9 | */ 10 | public class HammingSimilarity extends ClassicSimilarity { 11 | 12 | @Override 13 | public float tf(float freq) { 14 | 15 | if(freq > 0) 16 | { 17 | return 1; 18 | } 19 | else { 20 | return 0; 21 | } 22 | } 23 | 24 | @Override 25 | public float lengthNorm(int length) 26 | { 27 | return 1; 28 | } 29 | 30 | @Override 31 | public float sloppyFreq(int distance) 32 | { 33 | return 1.0f; 34 | } 35 | 36 | @Override 37 | public float idf(long docFreq, long numDocs) 38 | { 39 | return 1.0f; 40 | } 41 | } 42 | -------------------------------------------------------------------------------- /solr_configs/solrconfig_configure_payloadedismax_parser.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 12 | 13 | 16 | 17 | 18 | 19 | 20 | 21 | -------------------------------------------------------------------------------- /python/lsh_variants.py: -------------------------------------------------------------------------------- 1 | import numpy 2 | 3 | def generate_random_vector(shape): 4 | v = numpy.random.normal(loc=0.0, scale=0.2, size=shape) 5 | l = numpy.linalg.norm(v) 6 | return v/l 7 | 8 | class LSHHasher(object): 9 | def __init__(self, num_vectors, vector_shape): 10 | self.num_vectors = num_vectors 11 | self.vector_shape = vector_shape 12 | self.random_vectors = [generate_random_vector(self.vector_shape) for i in range(self.num_vectors)] 13 | 14 | def hash_vector(self, vector, num_bits=None, as_str=False): 15 | if num_bits is None: 16 | num_bits = self.num_vectors 17 | assert num_bits <= self.num_vectors, "Can't have more bits than vectors" 18 | bits = [] 19 | for random_vec in self.random_vectors[:num_bits]: 20 | cos_sim = numpy.dot(vector, random_vec) 21 | hash_bit = +1 if cos_sim >= 0 else -1 22 | if as_str: 23 | hash_bit = "+1" if hash_bit == 1 else "-1" 24 | bits.append(hash_bit) 25 | return bits -------------------------------------------------------------------------------- /solr_plugins/src/main/java/org/dice/solrenhancements/queryparsers/PayloadAwareExtendedDismaxQParserPlugin.java: -------------------------------------------------------------------------------- 1 | package org.dice.solrenhancements.queryparsers; 2 | 3 | import org.dice.solrenhancements.JarVersion; 4 | import org.slf4j.Logger; 5 | import org.slf4j.LoggerFactory; 6 | 7 | /** 8 | * Created by simon.hughes on 3/29/14. 9 | */ 10 | public class PayloadAwareExtendedDismaxQParserPlugin extends org.apache.solr.search.ExtendedDismaxQParserPlugin { 11 | public static final java.lang.String NAME = "payloadEdismax"; 12 | 13 | public PayloadAwareExtendedDismaxQParserPlugin() { 14 | super(); 15 | } 16 | 17 | public org.apache.solr.search.QParser createParser( 18 | java.lang.String qstr, 19 | org.apache.solr.common.params.SolrParams localParams, 20 | org.apache.solr.common.params.SolrParams params, 21 | org.apache.solr.request.SolrQueryRequest req) 22 | { 23 | return new PayloadAwareExtendDismaxQParser(qstr, localParams, params, req); 24 | } 25 | 26 | private static final Logger Log = LoggerFactory.getLogger(PayloadAwareExtendedDismaxQParserPlugin.class); 27 | 28 | private String version = null; 29 | 30 | } -------------------------------------------------------------------------------- /solr_plugins/src/main/java/org/dice/solrenhancements/similarity/PayloadOnlySimilarity.java: -------------------------------------------------------------------------------- 1 | package org.dice.solrenhancements.similarity; 2 | 3 | import org.apache.lucene.analysis.payloads.PayloadHelper; 4 | import org.apache.lucene.util.BytesRef; 5 | 6 | import javax.swing.plaf.DesktopIconUI; 7 | 8 | /** 9 | * Created by simon.hughes on 4/16/14. 10 | */ 11 | public class PayloadOnlySimilarity extends DiceDefaultSimilarity { 12 | 13 | @Override 14 | public float sloppyFreq(int distance) 15 | { 16 | return 1.0f; 17 | } 18 | 19 | @Override 20 | public float tf(float freq) { 21 | 22 | if(freq > 0){ 23 | return 1.0f; 24 | } 25 | return 0.0f; 26 | } 27 | 28 | @Override 29 | public float idf(long docFreq, long numDocs) 30 | { 31 | return 1.0f; 32 | } 33 | 34 | @Override 35 | public float lengthNorm(int length) 36 | { 37 | return 1; 38 | } 39 | 40 | @Override 41 | public float scorePayload(int doc, int start, int end, BytesRef payload) { 42 | if (payload != null) { 43 | float x = PayloadHelper.decodeFloat(payload.bytes, payload.offset); 44 | return x; 45 | } 46 | return 1.0F; 47 | } 48 | 49 | } 50 | -------------------------------------------------------------------------------- /solr_plugins/src/main/java/org/dice/solrenhancements/JarVersion.java: -------------------------------------------------------------------------------- 1 | package org.dice.solrenhancements; 2 | 3 | import org.slf4j.Logger; 4 | 5 | import java.io.InputStream; 6 | import java.net.URL; 7 | import java.util.Enumeration; 8 | 9 | /** 10 | * Created by simon.hughes on 7/7/16. 11 | */ 12 | public class JarVersion { 13 | 14 | private class stub{ 15 | 16 | } 17 | 18 | public static String getVersion(Logger log){ 19 | 20 | Enumeration resources; 21 | StringBuilder stringBuilder = new StringBuilder(); 22 | 23 | try { 24 | resources = stub.class.getClassLoader().getResources("META-INF/MANIFEST.MF"); 25 | while (resources.hasMoreElements()) { 26 | URL url = resources.nextElement(); 27 | /* let's not read other jar's manifests */ 28 | if (!url.toString().contains("DiceSolrEnhancements")) { 29 | continue; 30 | } 31 | InputStream reader = url.openStream(); 32 | while(reader.available() > 0) { 33 | char c = (char) reader.read(); 34 | stringBuilder.append(c); 35 | /* skip lines that don't contain the built-date */ 36 | if (stringBuilder.toString().contains(System.getProperty("line.separator")) && 37 | !stringBuilder.toString().contains("Build-Time")) { 38 | stringBuilder.setLength(0); 39 | } 40 | } 41 | } 42 | } catch (Exception e) { 43 | log.warn("Failed to read manifest during request for version!"); 44 | return "Error reading manifest!"; 45 | } 46 | return stringBuilder.toString(); 47 | } 48 | } 49 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | db.sqlite3 58 | 59 | # Flask stuff: 60 | instance/ 61 | .webassets-cache 62 | 63 | # Scrapy stuff: 64 | .scrapy 65 | 66 | # Sphinx documentation 67 | docs/_build/ 68 | 69 | # PyBuilder 70 | target/ 71 | 72 | # Jupyter Notebook 73 | .ipynb_checkpoints 74 | 75 | # pyenv 76 | .python-version 77 | 78 | # celery beat schedule file 79 | celerybeat-schedule 80 | 81 | # SageMath parsed files 82 | *.sage.py 83 | 84 | # Environments 85 | .env 86 | .venv 87 | env/ 88 | venv/ 89 | ENV/ 90 | env.bak/ 91 | venv.bak/ 92 | 93 | # Spyder project settings 94 | .spyderproject 95 | .spyproject 96 | 97 | # Rope project settings 98 | .ropeproject 99 | 100 | # mkdocs documentation 101 | /site 102 | 103 | # mypy 104 | .mypy_cache/ 105 | 106 | # java 107 | *.class 108 | *.iml 109 | 110 | model 111 | -------------------------------------------------------------------------------- /python/vector_thresholding.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from numpy.linalg import norm 3 | 4 | def vec2unit_vec(vec): 5 | return vec / norm(vec) 6 | 7 | # thresholds values based on percentiles 8 | # so values above get set to pos_val, values below get set to neg_val, and in between get set to middle_val 9 | def threshold_vector_by_pct(vector, pct_cutoff=90, pos_val=1, neg_val=-1, middle_val=0): 10 | assert pct_cutoff >= 50 11 | neg_pct_cutoff = 100-pct_cutoff 12 | 13 | pos_threshold = np.percentile(vector, pct_cutoff) 14 | neg_threshold = np.percentile(vector, neg_pct_cutoff) 15 | 16 | mod_vec = vector 17 | mod_vec = np.where( mod_vec >= pos_threshold, pos_val, 18 | np.where(mod_vec < neg_threshold, neg_val, 19 | middle_val)) 20 | return mod_vec 21 | 22 | def sparsify_vector_by_pct(vector, pct_cutoff=90): 23 | assert pct_cutoff >= 50 24 | neg_pct_cutoff = 100-pct_cutoff 25 | 26 | pos_threshold = np.percentile(vector, pct_cutoff) 27 | neg_threshold = np.percentile(vector, neg_pct_cutoff) 28 | 29 | mod_vec = vector.copy() 30 | mod_vec[(mod_vec <= pos_threshold) & (mod_vec >= neg_threshold)] = 0 31 | return mod_vec 32 | 33 | # thresholds values based on pct's computed from a population of values 34 | def threshold_vector_by_popn_pct(vector, pct2val, pct_cutoff=90, pos_val=1, neg_val=-1, middle_val=0): 35 | assert pct_cutoff >= 50 36 | neg_pct_cutoff = 100-pct_cutoff 37 | mod_vec = vector 38 | mod_vec = np.where( mod_vec >= pct2val[pct_cutoff], pos_val, 39 | np.where(mod_vec < pct2val[neg_pct_cutoff], neg_val, 40 | middle_val)) 41 | return mod_vec 42 | 43 | # thresholds values to above (inc.) and below a threshold 44 | def threshold_vector_by_val(vector, cutoff=0, pos_val=1, neg_val=-1): 45 | mod_vec = vector 46 | mod_vec = np.where( mod_vec >= cutoff, pos_val, neg_val) 47 | return mod_vec -------------------------------------------------------------------------------- /solr_configs/schema_xml_field_definitions.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | cluster 34 | 35 | 36 | -------------------------------------------------------------------------------- /python/solr_encode_vectors.py: -------------------------------------------------------------------------------- 1 | # Indexing functions 2 | def solr_encode_vector(vector): 3 | tokens = ["{i}|{val}".format(i=i, val=val) for i, val in enumerate(vector)] 4 | return " ".join(tokens) 5 | 6 | def solr_encode_sparse_vector(vector): 7 | tokens = ["{i}|{val}".format(i=i, val=val) for i, val in enumerate(vector) if val != 0.0] 8 | return " ".join(tokens) 9 | 10 | # encodes sparse vector as tokens, including weights 11 | def solr_encode_quantize_sparse_vector(vector, decimal_places=2): 12 | tokens = ["{i}_{dp}dp_{sign}_{val}".format(i=i, dp=decimal_places, 13 | sign="neg" if val < 0 else "pos", 14 | val=round(abs(val),decimal_places)) 15 | for i, val in enumerate(vector) if val != 0.0] 16 | 17 | return tokens 18 | 19 | def solr_encode_hash_finger_print(hash_vector): 20 | return ["".join(map(lambda s: str(s).rjust(2, '+'), hash_vector))] 21 | 22 | # Query functions 23 | def __tokens_to_field_query__(field, stokens): 24 | return "{field}:({stokens})".format(field=field, stokens=stokens) 25 | 26 | def __build_field_query__(field, vector, sparse=True): 27 | stokens = " ".join(["{i}^{val}".format(field=field, i=i, val=val) 28 | for i, val in enumerate(vector) if (sparse and val != 0.0) or not sparse]) 29 | return __tokens_to_field_query__(field, stokens) 30 | 31 | def solr_encode_vector_for_query(vector, field): 32 | return __build_field_query__(field, vector, sparse=False) 33 | 34 | def solr_encode_sparse_vector_for_query(vector, field): 35 | return __build_field_query__(field, vector, sparse=True) 36 | 37 | def __tokenize_vector_component__(i, val, decimal_places): 38 | rounded_val = round(abs(val),decimal_places) 39 | return "{i}_{dp}dp_{sign}_{val}".format(i=i, dp=decimal_places, 40 | sign="neg" if val < 0 else "pos", 41 | val=rounded_val) 42 | 43 | # encodes sparse vector as tokens, including weights 44 | def solr_encode_quantize_sparse_vector_for_query(vector, field, decimal_places=2): 45 | stokens = " ".join(["{i}_{dp}dp_{sign}_{val}".format(i=i, dp=decimal_places, 46 | sign="neg" if val < 0 else "pos", 47 | val=round(abs(val),decimal_places)) 48 | for i, val in enumerate(vector) if val != 0.0]) 49 | 50 | return __tokens_to_field_query__(field, stokens) 51 | 52 | def solr_encode_hash_finger_print_for_query(hash_vector, field): 53 | finger_prints = solr_encode_hash_finger_print(hash_vector) 54 | assert len(finger_prints) == 1, "Should only be one finger print" 55 | return __tokens_to_field_query__(field, "\"" + finger_prints[0] + "\"") 56 | -------------------------------------------------------------------------------- /solr_plugins/src/main/java/org/dice/solrenhancements/similarity/DiceDefaultSimilarity.java: -------------------------------------------------------------------------------- 1 | package org.dice.solrenhancements.similarity; 2 | 3 | /** 4 | * Created by simon.hughes on 6/3/15. 5 | */ 6 | /* 7 | * Licensed to the Apache Software Foundation (ASF) under one or more 8 | * contributor license agreements. See the NOTICE file distributed with 9 | * this work for additional information regarding copyright ownership. 10 | * The ASF licenses this file to You under the Apache License, Version 2.0 11 | * (the "License"); you may not use this file except in compliance with 12 | * the License. You may obtain a copy of the License at 13 | * 14 | * http://www.apache.org/licenses/LICENSE-2.0 15 | * 16 | * Unless required by applicable law or agreed to in writing, software 17 | * distributed under the License is distributed on an "AS IS" BASIS, 18 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 19 | * See the License for the specific language governing permissions and 20 | * limitations under the License. 21 | */ 22 | 23 | import org.apache.lucene.index.FieldInvertState; 24 | import org.apache.lucene.search.similarities.Similarity; 25 | import org.apache.lucene.search.similarities.TFIDFSimilarity; 26 | import org.apache.lucene.util.BytesRef; 27 | import org.apache.lucene.util.SmallFloat; 28 | 29 | public class DiceDefaultSimilarity extends TFIDFSimilarity { 30 | 31 | /** Cache of decoded bytes. */ 32 | private static final float[] NORM_TABLE = new float[256]; 33 | 34 | static { 35 | for (int i = 0; i < 256; i++) { 36 | NORM_TABLE[i] = SmallFloat.byte315ToFloat((byte)i); 37 | } 38 | } 39 | 40 | /** Sole constructor: parameter-free */ 41 | public DiceDefaultSimilarity() {} 42 | 43 | /** Implemented as sqrt(freq). */ 44 | @Override 45 | public float tf(float freq) { 46 | return (float)Math.sqrt(freq); 47 | } 48 | 49 | /** Implemented as 1 / (distance + 1). */ 50 | @Override 51 | public float sloppyFreq(int distance) { 52 | return 1.0f / (distance + 1); 53 | } 54 | 55 | /** The default implementation returns 1 */ 56 | @Override 57 | public float scorePayload(int doc, int start, int end, BytesRef payload) { 58 | return 1; 59 | } 60 | 61 | /** Implemented as log(numDocs/(docFreq+1)) + 1. */ 62 | @Override 63 | public float idf(long docFreq, long numDocs) { 64 | return (float)(Math.log(numDocs/(double)(docFreq+1)) + 1.0); 65 | } 66 | 67 | @Override 68 | public float lengthNorm(int length) { 69 | return (float)(1.0 / Math.sqrt(length)); 70 | } 71 | 72 | /** 73 | * True if overlap tokens (tokens with a position of increment of zero) are 74 | * discounted from the document's length. 75 | */ 76 | protected boolean discountOverlaps = true; 77 | 78 | /** Determines whether overlap tokens (Tokens with 79 | * 0 position increment) are ignored when computing 80 | * norm. By default this is true, meaning overlap 81 | * tokens do not count when computing norms. 82 | * 83 | * @lucene.experimental 84 | * 85 | * @see #computeNorm 86 | */ 87 | public void setDiscountOverlaps(boolean v) { 88 | discountOverlaps = v; 89 | } 90 | 91 | /** 92 | * Returns true if overlap tokens are discounted from the document's length. 93 | * @see #setDiscountOverlaps 94 | */ 95 | public boolean getDiscountOverlaps() { 96 | return discountOverlaps; 97 | } 98 | 99 | @Override 100 | public String toString() { 101 | return "DefaultSimilarity"; 102 | } 103 | } -------------------------------------------------------------------------------- /solr_plugins/src/main/java/org/dice/solrenhancements/queryparsers/PayloadAwareExtendDismaxQParser.java: -------------------------------------------------------------------------------- 1 | package org.dice.solrenhancements.queryparsers; 2 | 3 | import org.apache.lucene.analysis.Analyzer; 4 | import org.apache.lucene.index.Term; 5 | import org.apache.lucene.queries.payloads.PayloadDecoder; 6 | import org.apache.lucene.queries.payloads.PayloadFunction; 7 | import org.apache.lucene.queries.payloads.PayloadScoreQuery; 8 | import org.apache.lucene.search.BoostQuery; 9 | import org.apache.lucene.search.Query; 10 | import org.apache.lucene.queries.payloads.AveragePayloadFunction; 11 | import org.apache.lucene.search.spans.SpanQuery; 12 | import org.apache.lucene.search.spans.SpanTermQuery; 13 | import org.apache.solr.schema.SchemaField; 14 | import org.apache.solr.search.ExtendedDismaxQParser; 15 | import org.apache.solr.search.SyntaxError; 16 | 17 | import java.util.HashMap; 18 | 19 | /** 20 | * Created by simon.hughes on 3/29/14. 21 | */ 22 | 23 | public class PayloadAwareExtendDismaxQParser extends ExtendedDismaxQParser { 24 | 25 | public PayloadAwareExtendDismaxQParser( 26 | java.lang.String qstr, 27 | org.apache.solr.common.params.SolrParams localParams, 28 | org.apache.solr.common.params.SolrParams params, 29 | org.apache.solr.request.SolrQueryRequest req) 30 | { 31 | super(qstr,localParams, params, req); 32 | } 33 | 34 | @Override 35 | protected org.apache.solr.search.ExtendedDismaxQParser.ExtendedSolrQueryParser createEdismaxQueryParser(org.apache.solr.search.QParser qParser, java.lang.String field) 36 | { 37 | return new PayloadAwareExtendedSolrQueryParser(qParser, field); 38 | } 39 | 40 | public static class PayloadAwareExtendedSolrQueryParser extends ExtendedDismaxQParser.ExtendedSolrQueryParser { 41 | 42 | public PayloadAwareExtendedSolrQueryParser(org.apache.solr.search.QParser parser, java.lang.String defaultField) { 43 | super(parser, defaultField); 44 | } 45 | 46 | @Override 47 | protected Query getFieldQuery(String field, String queryText, boolean quoted) throws SyntaxError { 48 | SchemaField sf = this.schema.getFieldOrNull(field); 49 | 50 | //TODO cache this check 51 | if (sf != null) { 52 | 53 | final String fieldTypeName = sf.getType().getTypeName().toLowerCase(); 54 | if(fieldTypeName.contains("payload") || fieldTypeName.contains("vector")) { 55 | // We need includeSpanScore to include the boost value also 56 | return new PayloadScoreQuery(new SpanTermQuery(new Term(field, queryText)), new AveragePayloadFunction(), null, true); 57 | } 58 | } 59 | return super.getFieldQuery(field, queryText, quoted); 60 | } 61 | 62 | @Override 63 | protected Query newFieldQuery(Analyzer analyzer, String field, String queryText, 64 | boolean quoted, boolean fieldAutoGenPhraseQueries, boolean enableGraphQueries, 65 | SynonymQueryStyle synonymQueryStyle) 66 | throws SyntaxError { 67 | Analyzer actualAnalyzer = parser.getReq().getSchema().getFieldType(field).getQueryAnalyzer(); 68 | SchemaField sf = this.schema.getFieldOrNull(field); 69 | if (sf != null) { 70 | 71 | final String fieldTypeName = sf.getType().getTypeName().toLowerCase(); 72 | if(fieldTypeName.contains("payload") || fieldTypeName.contains("vector")) { 73 | // We need includeSpanScore to include the boost value also 74 | return new PayloadScoreQuery(new SpanTermQuery(new Term(field, queryText)), new AveragePayloadFunction(), null, true); 75 | } 76 | } 77 | 78 | return super.newFieldQuery(actualAnalyzer, field, queryText, quoted, fieldAutoGenPhraseQueries, enableGraphQueries, synonymQueryStyle); 79 | } 80 | } 81 | 82 | } 83 | -------------------------------------------------------------------------------- /Readme.md: -------------------------------------------------------------------------------- 1 | # Vectors in Search 2 | 3 | Dice.com code for implementing the ideas discussed in the following talks: 4 | 5 | * 'Vectors in Search' - [Activate 2018 conference](https://activate-conf.com/more-events/) 6 | * 'Searching with Vectors' - [Haystack 2019 conference](https://haystackconf.com/2019/vectors/) 7 | 8 | This extends my earlier work on 'Conceptual Search' which can be found here - https://github.com/DiceTechJobs/ConceptualSearch (including slides and video links). In this talk, I present a number of different approaches for searching vectors at scale using an inverted index. This implements approaches to [Approximate k-Nearest Neighbor Search](https://en.wikipedia.org/wiki/Nearest_neighbor_search#Approximate_nearest_neighbor) including: 9 | 10 | - LSH (using the Sim Hash) 11 | - K-Means Tree 12 | - Vector Thresholding 13 | 14 | and describes how these ideas can be implemented and queried efficiently within an inverted index. 15 | 16 | **UPDATE:** 17 | After talking with [Trey Grainger](https://www.linkedin.com/in/treygrainger/) and [Erik Hatcher](https://www.linkedin.com/in/erik-hatcher-94820/) from LucidWorks, they recommended using term frequency in place of payloads for the solutions where I embed term weights into the index and use a special payload aware similarity function (which would also not be needed). Payloads incur a significant performance penalty. The challenge with this is the negative weights, I assume it is not possible to encode negative term frequencies, but this can be worked around by having different tokens for positive and negative weighted tokens, and making similar adjustments at query time (where negative boosts can be applied in Solr as needed). 18 | 19 | Lucene Documentation: [Lucene Delimited Term Frequency Filter](https://lucene.apache.org/core/7_0_0/analyzers-common/org/apache/lucene/analysis/miscellaneous/DelimitedTermFrequencyTokenFilter.html) 20 | 21 | There has also been a recent update to Lucene core that is applicable here and is soon to make it's way into Elastic search at time of writing: [Block Max WAND](https://www.elastic.co/blog/faster-retrieval-of-top-hits-in-elasticsearch-with-block-max-wand). This produces a signifcant speed up for large boolean OR queries where you don't need to know the exact number of results but just care about getting the top-N results as fast as possible. All of the approaches I discuss here generate relatively large OR queries and so this is very relevant. I have also read that the current implementation of minimum-should-match also includes similar optimizations, and so the same sort of performance gain may already be attained using appropriate mm settings, something that I was already experimenting with in my code. 22 | 23 | ## Directory Structure 24 | - **python** 25 | - Code for implementing the k-means tree, LSH sim hash and vector thresholding algorithms, and indexing and searching vectors in solr using these techniques. 26 | - **solr_plugins** 27 | - Java code for implementing the custom similarity classes and payloadEdismax parser described in the talk. 28 | - **solr_configs** 29 | - Xml snippets for importing the solr plugins from the 'solr_vectors_in_search_plugins' java code. 30 | 31 | ## Implementation Details 32 | - Solr Version - 7.5 33 | - Python Version - 3.x+ (3.5 used) 34 | 35 | ## Links to Talks 36 | 37 | * **Activate 2018:** 'Vectors in Search' 38 | * [Slides](https://www.slideshare.net/lucidworks/vectors-in-search-towards-more-semantic-matching-simon-hughes-dicecom?qid=4c9af9c0-0554-4251-bd47-9345ff508569&v=&b=&from_search=2) 39 | * [Video](https://www.youtube.com/watch?v=rSDqhGn_8Zo&list=PLU6n9Voqu_1HW8-VavVMa9lP8-oF8Oh5t&index=21&t=56s) 40 | 41 | * **Haystack 2019:** 'Searching with Vectors' 42 | * [Slides](https://www.slideshare.net/o19s/haystack-2019-search-with-vectors-simon-hughes) 43 | * [Video](https://www.youtube.com/watch?v=hycH6Rn4RaU&list=PLCoJWKqBHERu9Fe0W12D7XKwGT2eoJJNU&index=19) 44 | 45 | ## Author 46 | Simon Hughes ( Chief Data Scientist, Dice.com ) 47 | * LinkedIn - https://www.linkedin.com/in/simon-hughes-data-scientist/ 48 | * Twitter - https://twitter.com/hughes_meister -------------------------------------------------------------------------------- /solr_plugins/pom.xml: -------------------------------------------------------------------------------- 1 | 2 | 5 | 4.0.0 6 | 7 | 8 | ${maven.build.timestamp} 9 | yyyyMMddHHmmss 10 | 11 | 12 | com.dice.solr.plugins 13 | DiceSolrEnhancements 14 | 15 | 1.0 16 | 17 | 18 | 19 | 20 | 21 | com.google.guava 22 | guava 23 | 12.0 24 | 25 | 26 | 27 | 28 | org.apache.solr 29 | solr-core 30 | 7.5.0 31 | 32 | 33 | 34 | org.apache.solr 35 | solr-solrj 36 | 7.5.0 37 | 38 | 39 | 40 | 41 | org.apache.lucene 42 | lucene-analyzers-common 43 | 7.5.0 44 | 45 | 46 | org.apache.lucene 47 | lucene-queryparser 48 | 7.5.0 49 | 50 | 51 | org.apache.lucene 52 | lucene-queries 53 | 7.5.0 54 | 55 | 56 | org.apache.lucene 57 | lucene-core 58 | 7.5.0 59 | 60 | 61 | org.json 62 | json 63 | 20131018 64 | 65 | 66 | 67 | junit 68 | junit 69 | 4.11 70 | 71 | 72 | 73 | org.apache.maven.plugins 74 | maven-deploy-plugin 75 | 2.8.2 76 | 77 | 78 | 79 | 80 | 81 | 82 | 83 | org.apache.maven.plugins 84 | maven-jar-plugin 85 | 2.4 86 | 87 | true 88 | 89 | 90 | 91 | 92 | org.apache.maven.plugins 93 | maven-antrun-plugin 94 | 95 | 96 | generate-resources 97 | 98 | run 99 | 100 | 101 | 102 | 103 | 104 | 105 | 106 | 107 | 108 | 109 | 110 | 111 | 112 | org.apache.maven.plugins 113 | maven-compiler-plugin 114 | 3.0 115 | 116 | 1.6 117 | 1.6 118 | 119 | 120 | 121 | 122 | 123 | 124 | src/main/resources 125 | true 126 | 127 | 128 | 129 | 130 | 131 | 132 | 133 | release 134 | 1fb7a15dd864-releases 135 | http://artifactory.services.dicedev.dhiaws.com/artifactory/libs-release-local 136 | 137 | 138 | 139 | -------------------------------------------------------------------------------- /python/kmeans_tree.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from collections import defaultdict 3 | from sklearn.cluster import KMeans 4 | 5 | class Labeller(object): 6 | # get's new node labels 7 | def __init__(self): 8 | self.current = -1 # needs to be -1 so first one is 0 9 | 10 | def get_new_node_id(self): 11 | self.current += 1 12 | return str(self.current) 13 | 14 | class KMeansTree(object): 15 | ROOT = "ROOT" 16 | 17 | def __init__(self, branch_factor, n_init=15, max_iter=300, n_jobs=1, max_cluster_size=None, verbose=False): 18 | self.branch_factor = branch_factor 19 | self.lbl = Labeller() 20 | self.id2kmeans = dict() # maps node-id to the kmeans cluster that created it's childen 21 | self.id2centroid = dict() 22 | self.n_init = n_init 23 | self.max_iter = max_iter 24 | self.tree = {} 25 | self.nodes = [] 26 | self.vectors = [] 27 | # maps vector indices to cluster node_ids 28 | self.ix2leaf_node_id = dict() 29 | self.leaf_nodeid2ixs = dict() 30 | self.depth_to_node_id = defaultdict(set) 31 | self.depth_to_node_id[0].add(KMeansTree.ROOT) 32 | self.id2_depth = dict() 33 | self.id2_depth[KMeansTree.ROOT] = 0 34 | self.verbose = verbose 35 | # for fast sub-tree querying 36 | self.id2sub_tree = dict() 37 | self.max_cluster_size = max_cluster_size if max_cluster_size else self.branch_factor 38 | assert self.max_cluster_size >= self.branch_factor, "Max cluster size must be >= branch_factor" 39 | # debugging 40 | self.max_depth = -1 41 | 42 | def is_a_leaf_node(self, node_id): 43 | return node_id in self.leaf_nodeid2ixs 44 | 45 | def get_subtree_for_node(self, node_id): 46 | return self.id2sub_tree[node_id] 47 | 48 | def get_leaf_node_for_ix(self, ix): 49 | assert ix in self.ix2leaf_node_id, "Index {ix} not found".format(ix=ix) 50 | return self.ix2leaf_node_id[ix] 51 | 52 | # get's all child indices for a node 53 | def get_indices_for_node(self, node_id): 54 | if self.is_a_leaf_node(node_id): 55 | return self.leaf_nodeid2ixs[node_id] 56 | else: 57 | indices = [] 58 | for node_id, _ in self.id2sub_tree[node_id].items(): 59 | indices.extend(self.get_indices_for_node(node_id)) 60 | return indices 61 | 62 | def get_leaf_nodes(self): 63 | return list(self.leaf_nodeid2ixs.keys()) 64 | 65 | def get_internal_nodes(self): 66 | return [node_id for node_id in self.nodes if not self.is_a_leaf_node(node_id)] 67 | 68 | def get_nodes_at_level(self, level): 69 | return self.depth_to_node_id[level] 70 | 71 | def __add_leaf_node__(self, ixs, parent_node_id, depth): 72 | if self.verbose: 73 | print("{indent}L - depth={depth}, size={size}, parent={parent_node_id}".format( 74 | indent="\t" * depth, size=len(ixs), depth=depth, parent_node_id=parent_node_id)) 75 | self.leaf_nodeid2ixs[parent_node_id] = ixs 76 | for ix in ixs: 77 | assert ix not in self.ix2leaf_node_id, "Index already mapped" 78 | self.ix2leaf_node_id[ix] = parent_node_id 79 | self.depth_to_node_id[depth].add(parent_node_id) 80 | return dict() # return empty dictionary as no children 81 | 82 | def __build_tree__(self, vecs, ixs, parent_node_id, depth): 83 | 84 | if self.verbose and depth > self.max_depth: 85 | print("New Max Depth={depth}, size={size}".format(depth=depth, size=len(vecs))) 86 | self.max_depth = max(self.max_depth, depth) 87 | 88 | if len(vecs) <= self.max_cluster_size: 89 | return self.__add_leaf_node__(ixs=ixs, parent_node_id=parent_node_id, depth=depth) 90 | 91 | if self.verbose: 92 | print("{indent}I - depth={depth}, size: {size}, parent: {parent_node_id}".format( 93 | indent="\t" * depth, size=len(vecs), depth=depth, parent_node_id=parent_node_id)) 94 | # TODO - get node id, make recursive, store depth 95 | km = KMeans(n_clusters=min(self.branch_factor, len(vecs)), n_init=self.n_init, max_iter=self.max_iter) 96 | km.fit(vecs) 97 | # store kmeans for later just in case 98 | self.id2kmeans[parent_node_id] = km 99 | 100 | assert len(km.labels_) == len(vecs), "|labels| != |vecs|" 101 | assert len(ixs) == len(vecs), "|indices| != |vecs|" 102 | 103 | # group vectors by cluster labels 104 | lbl2vecs = defaultdict(list) 105 | lbl2ixs = defaultdict(list) 106 | for lbl, ix, vec in zip(km.labels_, ixs, vecs): 107 | lbl2vecs[lbl].append(vec) 108 | lbl2ixs[lbl].append(ix) 109 | 110 | lbls = lbl2vecs.keys() 111 | if len(lbls) == 1: 112 | # if only one cluster found - items could all be identical (I have found instances of this) 113 | # make a leaf node 114 | return self.__add_leaf_node__(ixs=ixs, parent_node_id=parent_node_id, depth=depth) 115 | 116 | tree = {} 117 | for lbl, child_vecs in lbl2vecs.items(): 118 | child_ixs = lbl2ixs[lbl] 119 | child_node_id = "{parent_node_id}->{childid}".format(parent_node_id=parent_node_id, 120 | childid=self.lbl.get_new_node_id()) 121 | self.nodes.append(child_node_id) 122 | self.depth_to_node_id[depth + 1].add(child_node_id) 123 | centroid = km.cluster_centers_[lbl] 124 | # we need to normalize centroids so we can do np.dot 125 | norm_centroid = centroid / np.linalg.norm(centroid) 126 | self.id2centroid[child_node_id] = norm_centroid 127 | sub_tree = self.__build_tree__(vecs=child_vecs, ixs=child_ixs, parent_node_id=child_node_id, 128 | depth=depth + 1) 129 | self.id2sub_tree[child_node_id] = sub_tree 130 | tree[child_node_id] = sub_tree 131 | 132 | # build node_id to depth mapping 133 | for depth, node_ids in self.depth_to_node_id.items(): 134 | for node_id in node_ids: 135 | self.id2_depth[node_id] = depth 136 | return tree 137 | 138 | def __build_labels__(self): 139 | # we have to now map the initial vectors to their leaf clusters, for sklearn structure 140 | leaf_labeller = Labeller() 141 | labels_ = [] 142 | leaf_node_labels_ = [] # retain the original labelling scheme 143 | node2lbl = dict() 144 | for ix in sorted(self.ix2leaf_node_id.keys()): # sort just in case hash ordering changes 145 | parent_node_id = self.ix2leaf_node_id[ix] 146 | leaf_node_labels_.append(parent_node_id) 147 | 148 | if parent_node_id in node2lbl: 149 | labels_.append(node2lbl[parent_node_id]) 150 | else: 151 | new_lbl = int(leaf_labeller.get_new_node_id()) 152 | labels_.append(new_lbl) 153 | node2lbl[parent_node_id] = new_lbl 154 | return labels_, leaf_node_labels_ 155 | 156 | def fit(self, vecs): 157 | # get the ixs of the vecs so we can return the labels latter 158 | self.vectors = vecs 159 | ixs = np.arange(len(vecs)) 160 | self.tree[KMeansTree.ROOT] = self.__build_tree__(vecs=vecs, ixs=ixs, parent_node_id=KMeansTree.ROOT, depth=0) 161 | # make sure ROOT is mapped 162 | self.id2sub_tree[KMeansTree.ROOT] = self.tree[KMeansTree.ROOT] 163 | self.labels_, self.leaf_node_labels_ = self.__build_labels__() 164 | assert len(self.labels_) == len(vecs), "|labels| != |vectors|" 165 | -------------------------------------------------------------------------------- /python/searchers.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from queue import PriorityQueue 3 | from kmeans_tree import KMeansTree 4 | 5 | class BruteForceSearcher(object): 6 | """ 7 | Does a brute force k-nn search. This is our gold standard, we want to get close to this in terms of accuracy 8 | But be must faster 9 | """ 10 | 11 | def __init__(self, km_tree): 12 | self.km_tree = km_tree 13 | self.vecs = km_tree.vectors 14 | self.stacked_vecs = np.vstack(self.vecs) 15 | 16 | def search(self, vector, k_neighbors=10): 17 | # compute sims 18 | dot_prod = np.dot(vector.reshape(1, len(vector)), self.stacked_vecs.T) 19 | dot_prod = dot_prod[0, :] 20 | ixs = np.argsort(dot_prod)[::-1] # reverse order so in best to worse 21 | top_ixs = ixs[:k_neighbors] 22 | result = [] 23 | for ix in top_ixs: 24 | sim = dot_prod[ix] 25 | result.append((sim, ix)) 26 | return result 27 | 28 | class BestBinFirstSearcher(object): 29 | def __init__(self, km_tree, default_max_points_to_search=100): 30 | self.km_tree = km_tree 31 | self.default_max_points_to_search = default_max_points_to_search 32 | 33 | def __compute_similarity__(self, vec1, vec2): 34 | return np.dot(vec1, vec2) 35 | 36 | def __compute_similarity__(self, vec1, vec2): 37 | return np.dot(vec1, vec2) 38 | 39 | def __find_best_leaf_node__(self, parent_node_id, vector, q): 40 | # This can happen im the while loop if nodes pushed on the q are leaf nodes 41 | if self.km_tree.is_a_leaf_node(parent_node_id): 42 | leaf_centroid = self.km_tree.id2centroid[parent_node_id] 43 | cosine_sim = self.__compute_similarity__(vector, leaf_centroid) 44 | # print("Best node: {node} with sim: {sim}".format(node=parent_node_id, sim=cosine_sim)) 45 | return -cosine_sim, parent_node_id 46 | 47 | sub_tree = self.km_tree.get_subtree_for_node(parent_node_id) 48 | if len(sub_tree) == 0: 49 | raise Exception("Sub tree should not be empty") 50 | 51 | tmp_q = PriorityQueue() 52 | for child_node_id, child_sub_tree in sub_tree.items(): 53 | centroid = self.km_tree.id2centroid[child_node_id] 54 | cosine_sim = self.__compute_similarity__(vector, centroid) 55 | # NOTE - this is a min heap, so returns the lowest value, so negate 56 | tmp_q.put((-cosine_sim, child_node_id)) 57 | 58 | # pop the top sim and node from the tmp priority q 59 | best_sim, best_node = tmp_q.get() 60 | # we need to add the items here so they don't include the current best node 61 | # which we just popped and removed - we only want the remaining nodes 62 | while not tmp_q.empty(): 63 | # the get above removes the top item, which is the only reason this part of the code exists 64 | q.put(tmp_q.get()) 65 | 66 | if self.km_tree.is_a_leaf_node(best_node): 67 | # negative as it was negated when added to the queue 68 | # print("Best node: {node} with sim: {sim}".format(node=best_node, sim=best_sim)) 69 | return best_sim, best_node 70 | else: 71 | return self.__find_best_leaf_node__(best_node, vector, q) 72 | 73 | def __add_neighbors_from_leaf_node_to_q__(self, leaf_node_id, vector, result_ix_q): 74 | ixs = self.km_tree.leaf_nodeid2ixs[leaf_node_id] 75 | for ix in ixs: 76 | cosine_sim = self.__compute_similarity__(vector, self.km_tree.vectors[ix]) 77 | result_ix_q.put((-cosine_sim, ix)) 78 | return len(ixs) 79 | 80 | def search_best_leaf_nodes(self, vector, max_nodes_to_search=30, k_neighbors=None): 81 | if k_neighbors is not None: 82 | assert max_nodes_to_search >= k_neighbors, "Max Nodes must be >= k neighbors" 83 | # this contains -sim, node_id pairs 84 | q = PriorityQueue() 85 | nodes_searched = 0 86 | # this contains -sim, ix pairs, unlike q 87 | result_node_q = PriorityQueue() 88 | 89 | matching_docs = 0 90 | best_sim, best_node = self.__find_best_leaf_node__(KMeansTree.ROOT, vector, q) 91 | result_node_q.put((best_sim, best_node)) 92 | matching_docs += len(self.km_tree.get_indices_for_node(best_node)) 93 | nodes_searched += 1 94 | 95 | while not q.empty() and (nodes_searched < max_nodes_to_search or 96 | (k_neighbors is not None and matching_docs < k_neighbors)): 97 | _, next_best_point = q.get() 98 | best_sim, best_node = self.__find_best_leaf_node__(next_best_point, vector, q) 99 | result_node_q.put((best_sim, best_node)) 100 | matching_docs += len(self.km_tree.get_indices_for_node(best_node)) 101 | nodes_searched += 1 102 | 103 | best_nodes = [] 104 | while not result_node_q.empty(): 105 | sim, node = result_node_q.get() 106 | # reverse similarity 107 | best_nodes.append((-sim, node)) 108 | return best_nodes 109 | 110 | def search_subtree(self, vector): 111 | # for a vector, find the best leaf node, and then return all the nodes along that route, in sim order 112 | # this contains -sim, node_id pairs 113 | q = PriorityQueue() 114 | # this contains -sim, ix pairs, unlike q 115 | best_sim, best_node = self.__find_best_leaf_node__(KMeansTree.ROOT, vector, q) 116 | 117 | sub_tree_nodes = [] 118 | while not q.empty(): 119 | sim, node = q.get() 120 | # reverse similarity 121 | sub_tree_nodes.append((-sim, node)) 122 | return sub_tree_nodes 123 | 124 | def search(self, vector, k_neighbors=10, max_points_to_search=None): 125 | if max_points_to_search is None: 126 | max_points_to_search = self.default_max_points_to_search 127 | 128 | assert max_points_to_search >= k_neighbors, \ 129 | "Max Points to Search must be >= k neighbors. max={max}, k-neighbors={k_neighbors}".format( 130 | max=max_points_to_search, k_neighbors=k_neighbors 131 | ) 132 | # this contains -sim, node_id pairs 133 | q = PriorityQueue() 134 | points_searched = 0 135 | # this contains -sim, ix pairs, unlike q 136 | result_ix_q = PriorityQueue() 137 | 138 | best_sim, best_node = self.__find_best_leaf_node__(KMeansTree.ROOT, vector, q) 139 | 140 | num_added = self.__add_neighbors_from_leaf_node_to_q__(best_node, vector, result_ix_q) 141 | points_searched += num_added 142 | 143 | while not q.empty() and points_searched < max_points_to_search: 144 | _, next_best_point = q.get() 145 | _, best_node = self.__find_best_leaf_node__(next_best_point, vector, q) 146 | num_added = self.__add_neighbors_from_leaf_node_to_q__(best_node, vector, result_ix_q) 147 | points_searched += num_added 148 | 149 | k_best = [] 150 | while not result_ix_q.empty() and len(k_best) < k_neighbors: 151 | sim, ix = result_ix_q.get() 152 | # reverse similarity sign 153 | k_best.append((-sim, ix)) 154 | return k_best 155 | 156 | def __find_best_node__(self, parent_node_id, vector, q, depth, max_depth): 157 | 158 | # This can happen im the while loop if nodes pushed on the q are leaf nodes 159 | if self.km_tree.is_a_leaf_node(parent_node_id) or depth >= max_depth: 160 | centroid = self.km_tree.id2centroid[parent_node_id] 161 | cosine_sim = self.__compute_similarity__(vector, centroid) 162 | # print("Best node: {node} with sim: {sim}".format(node=parent_node_id, sim=cosine_sim)) 163 | return -cosine_sim, parent_node_id 164 | 165 | sub_tree = self.km_tree.get_subtree_for_node(parent_node_id) 166 | if len(sub_tree) == 0: 167 | raise Exception("Sub tree should not be empty") 168 | 169 | tmp_q = PriorityQueue() 170 | for child_node_id, child_sub_tree in sub_tree.items(): 171 | centroid = self.km_tree.id2centroid[child_node_id] 172 | cosine_sim = self.__compute_similarity__(vector, centroid) 173 | # NOTE - this is a min heap, so returns the lowest value, so negate 174 | tmp_q.put((-cosine_sim, child_node_id)) 175 | 176 | # pop (i.e. REMOVE) the top sim and node from the tmp priority q 177 | best_sim, best_node = tmp_q.get() 178 | # we need to add the items here so they don't include the current best node 179 | # which we just popped and removed - we only want the remaining nodes 180 | while not tmp_q.empty(): 181 | # the get above removes the top item, which is the only reason this part of the code exists 182 | q.put(tmp_q.get()) 183 | 184 | return self.__find_best_node__(best_node, vector, q, depth=depth+1, max_depth=max_depth) 185 | 186 | def search_best_nodes(self, vector, max_nodes_to_search=30, max_depth=None): 187 | if max_depth is None: 188 | max_depth = self.km_tree.max_depth 189 | assert max_depth > 0, "Max depth has to be at least 1" 190 | q = PriorityQueue() 191 | nodes_searched = 0 192 | # this contains -sim, ix pairs, unlike q 193 | result_node_q = PriorityQueue() 194 | 195 | matching_docs = 0 196 | best_sim, best_node = self.__find_best_node__(KMeansTree.ROOT, vector, q, depth=0, max_depth=max_depth) 197 | result_node_q.put((best_sim, best_node)) 198 | nodes_searched += 1 199 | 200 | while not q.empty() and nodes_searched < max_nodes_to_search : 201 | _, next_best_point = q.get() 202 | node_depth = self.km_tree.id2_depth[next_best_point] 203 | best_sim, best_node = self.__find_best_node__(next_best_point, vector, q, depth=node_depth, max_depth=max_depth) 204 | result_node_q.put((best_sim, best_node)) 205 | nodes_searched += 1 206 | 207 | best_nodes = [] 208 | while not result_node_q.empty(): 209 | sim, node = result_node_q.get() 210 | # reverse similarity 211 | best_nodes.append((-sim, node)) 212 | return best_nodes 213 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright [yyyy] [name of copyright owner] 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /python/Tutorial.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": {}, 6 | "source": [ 7 | "## Step 1 - Install Gensim, and Dowload the Word2Vec Model and vectors" 8 | ] 9 | }, 10 | { 11 | "cell_type": "code", 12 | "execution_count": null, 13 | "metadata": {}, 14 | "outputs": [], 15 | "source": [ 16 | "# Install gensim\n", 17 | "# !pip install gensim" 18 | ] 19 | }, 20 | { 21 | "cell_type": "code", 22 | "execution_count": 14, 23 | "metadata": {}, 24 | "outputs": [], 25 | "source": [ 26 | "import gensim\n", 27 | "import gensim.downloader as api" 28 | ] 29 | }, 30 | { 31 | "cell_type": "code", 32 | "execution_count": 28, 33 | "metadata": {}, 34 | "outputs": [ 35 | { 36 | "name": "stdout", 37 | "output_type": "stream", 38 | "text": [ 39 | "[==================================================] 100.0% 128.1/128.1MB downloaded\n" 40 | ] 41 | }, 42 | { 43 | "data": { 44 | "text/plain": [ 45 | "[('dog', 0.8798074722290039),\n", 46 | " ('rabbit', 0.7424427270889282),\n", 47 | " ('cats', 0.7323004007339478),\n", 48 | " ('monkey', 0.7288710474967957),\n", 49 | " ('pet', 0.7190139293670654),\n", 50 | " ('dogs', 0.7163873314857483),\n", 51 | " ('mouse', 0.6915251016616821),\n", 52 | " ('puppy', 0.6800068616867065),\n", 53 | " ('rat', 0.6641027331352234),\n", 54 | " ('spider', 0.6501134634017944)]" 55 | ] 56 | }, 57 | "execution_count": 28, 58 | "metadata": {}, 59 | "output_type": "execute_result" 60 | } 61 | ], 62 | "source": [ 63 | "model = api.load(\"glove-wiki-gigaword-100\") # download the model and return as object ready for use\n", 64 | "model.most_similar(\"cat\")" 65 | ] 66 | }, 67 | { 68 | "cell_type": "code", 69 | "execution_count": 41, 70 | "metadata": {}, 71 | "outputs": [ 72 | { 73 | "data": { 74 | "text/plain": [ 75 | "[('data', 0.7814346551895142),\n", 76 | " ('database', 0.6885414719581604),\n", 77 | " ('databases', 0.6605490446090698),\n", 78 | " ('knowledge', 0.6585058569908142),\n", 79 | " ('analysis', 0.6544734835624695),\n", 80 | " ('search', 0.6509679555892944),\n", 81 | " ('communication', 0.6452988386154175),\n", 82 | " ('documentation', 0.6334584951400757),\n", 83 | " ('processing', 0.6322544813156128),\n", 84 | " ('dissemination', 0.6290022134780884)]" 85 | ] 86 | }, 87 | "execution_count": 41, 88 | "metadata": {}, 89 | "output_type": "execute_result" 90 | } 91 | ], 92 | "source": [ 93 | "model.most_similar(positive=[\"information\",\"retrieval\"])" 94 | ] 95 | }, 96 | { 97 | "cell_type": "code", 98 | "execution_count": 39, 99 | "metadata": {}, 100 | "outputs": [ 101 | { 102 | "data": { 103 | "text/plain": [ 104 | "[('information', 0.6572738289833069),\n", 105 | " ('knowledge', 0.6555200815200806),\n", 106 | " ('human', 0.6344870328903198),\n", 107 | " ('biological', 0.6280955076217651),\n", 108 | " ('using', 0.6267763376235962),\n", 109 | " ('secret', 0.6181720495223999),\n", 110 | " ('use', 0.6163333654403687),\n", 111 | " ('scientific', 0.6116725206375122),\n", 112 | " ('communication', 0.6081548929214478),\n", 113 | " ('data', 0.6031312346458435)]" 114 | ] 115 | }, 116 | "execution_count": 39, 117 | "metadata": {}, 118 | "output_type": "execute_result" 119 | } 120 | ], 121 | "source": [ 122 | "model.most_similar(positive=[\"artificial\",\"intelligence\"])" 123 | ] 124 | }, 125 | { 126 | "cell_type": "code", 127 | "execution_count": 43, 128 | "metadata": {}, 129 | "outputs": [ 130 | { 131 | "data": { 132 | "text/plain": [ 133 | "[('the', ),\n", 134 | " (',', ),\n", 135 | " ('.', ),\n", 136 | " ('of', ),\n", 137 | " ('to', ),\n", 138 | " ('and', ),\n", 139 | " ('in', ),\n", 140 | " ('a', ),\n", 141 | " ('\"', ),\n", 142 | " (\"'s\", )]" 143 | ] 144 | }, 145 | "execution_count": 43, 146 | "metadata": {}, 147 | "output_type": "execute_result" 148 | } 149 | ], 150 | "source": [ 151 | "list(model.vocab.items())[0:10]" 152 | ] 153 | }, 154 | { 155 | "cell_type": "code", 156 | "execution_count": 48, 157 | "metadata": {}, 158 | "outputs": [ 159 | { 160 | "data": { 161 | "text/plain": [ 162 | "400000" 163 | ] 164 | }, 165 | "execution_count": 48, 166 | "metadata": {}, 167 | "output_type": "execute_result" 168 | } 169 | ], 170 | "source": [ 171 | "keys = model.vocab.keys()\n", 172 | "vecs = []\n", 173 | "key2vec = dict()\n", 174 | "for k in keys:\n", 175 | " key2vec[k] = model[k]\n", 176 | "len(key2vec)" 177 | ] 178 | }, 179 | { 180 | "cell_type": "code", 181 | "execution_count": 50, 182 | "metadata": {}, 183 | "outputs": [ 184 | { 185 | "name": "stdout", 186 | "output_type": "stream", 187 | "text": [ 188 | "the 5.8211536\n", 189 | ", 5.553375\n", 190 | ". 5.4601502\n", 191 | "of 6.296869\n", 192 | "to 6.450651\n", 193 | "and 5.667807\n", 194 | "in 6.0944986\n", 195 | "a 6.242884\n", 196 | "\" 6.5840864\n", 197 | "'s 6.662169\n" 198 | ] 199 | } 200 | ], 201 | "source": [ 202 | "import numpy\n", 203 | "from numpy.linalg import norm\n", 204 | "\n", 205 | "for k,v in list(key2vec.items())[:10]:\n", 206 | " print(k.ljust(10), norm(v))" 207 | ] 208 | }, 209 | { 210 | "cell_type": "code", 211 | "execution_count": 30, 212 | "metadata": {}, 213 | "outputs": [], 214 | "source": [ 215 | "# info = api.info() # show info about available models/datasets\n", 216 | "# from pprint import pprint\n", 217 | "# pprint(info)" 218 | ] 219 | }, 220 | { 221 | "cell_type": "code", 222 | "execution_count": 31, 223 | "metadata": {}, 224 | "outputs": [ 225 | { 226 | "data": { 227 | "text/plain": [ 228 | "[('nokesville', ),\n", 229 | " ('a&d', ),\n", 230 | " ('poucher', ),\n", 231 | " ('siwash', ),\n", 232 | " ('germanotta', ),\n", 233 | " ('crivitz', ),\n", 234 | " ('all-french', ),\n", 235 | " ('simrall', ),\n", 236 | " ('asperen', ),\n", 237 | " ('barshai', ),\n", 238 | " ('spider-slayer', ),\n", 239 | " ('rhins', ),\n", 240 | " ('1978-9', ),\n", 241 | " ('precrime', ),\n", 242 | " ('hbe', ),\n", 243 | " ('wls-fm', ),\n", 244 | " ('smartened', ),\n", 245 | " ('fourt', ),\n", 246 | " ('lilyan', ),\n", 247 | " ('molek', ),\n", 248 | " ('indal', ),\n", 249 | " ('bramford', ),\n", 250 | " ('israelitische', ),\n", 251 | " ('bivouacking', ),\n", 252 | " ('dabholkar', ),\n", 253 | " ('rosenholm', ),\n", 254 | " ('zonata', ),\n", 255 | " ('zoltar', ),\n", 256 | " ('qvt', ),\n", 257 | " ('dc-9-30', ),\n", 258 | " ('nettlebed', ),\n", 259 | " ('bosendorfer', ),\n", 260 | " ('sutovsky', ),\n", 261 | " ('nohab', ),\n", 262 | " ('veljovic', ),\n", 263 | " ('evar', ),\n", 264 | " ('stantonbury', ),\n", 265 | " ('ischgl', ),\n", 266 | " ('nyc-based', ),\n", 267 | " ('mpoyo', ),\n", 268 | " ('vuko', ),\n", 269 | " ('breznica', ),\n", 270 | " ('543-member', ),\n", 271 | " ('tanaji', ),\n", 272 | " ('smidgeon', ),\n", 273 | " ('pellam', ),\n", 274 | " ('stargrave', ),\n", 275 | " ('tessina', ),\n", 276 | " ('seehausen', ),\n", 277 | " ('manko', ),\n", 278 | " ('gun-brigs', ),\n", 279 | " ('in-the-fields', ),\n", 280 | " ('13-city', ),\n", 281 | " ('shiroma', ),\n", 282 | " ('detemobil', ),\n", 283 | " ('pondsmith', ),\n", 284 | " ('matovina', ),\n", 285 | " ('tuanjie', ),\n", 286 | " ('molinia', ),\n", 287 | " ('sneakiness', ),\n", 288 | " ('kaum', ),\n", 289 | " ('grindall', ),\n", 290 | " ('muv-luv', ),\n", 291 | " ('conne', ),\n", 292 | " ('stellman', ),\n", 293 | " ('rearviewmirror', ),\n", 294 | " ('smth', ),\n", 295 | " ('lelièvre', ),\n", 296 | " ('chelford', ),\n", 297 | " ('piesiewicz', ),\n", 298 | " ('sighişoara', ),\n", 299 | " ('arichat', ),\n", 300 | " ('florescence', ),\n", 301 | " ('pynn', ),\n", 302 | " ('eeu', ),\n", 303 | " ('dahomeyan', ),\n", 304 | " ('burgate', ),\n", 305 | " ('1i', ),\n", 306 | " ('marcognet', ),\n", 307 | " ('bouchout', ),\n", 308 | " ('kreiensen', ),\n", 309 | " ('cleistogamous', ),\n", 310 | " ('appleby-in-westmorland',\n", 311 | " ),\n", 312 | " ('songful', ),\n", 313 | " ('panchita', ),\n", 314 | " ('o’regan', ),\n", 315 | " ('brennabor', ),\n", 316 | " ('parleyed', ),\n", 317 | " ('sanderford', ),\n", 318 | " ('kargaly', ),\n", 319 | " ('jämsä', ),\n", 320 | " ('l’oréal', ),\n", 321 | " ('chamberlaine', ),\n", 322 | " ('bilanz', ),\n", 323 | " ('p166', ),\n", 324 | " ('harde', ),\n", 325 | " ('vru', ),\n", 326 | " ('subbuteo', ),\n", 327 | " ('amirav', ),\n", 328 | " ('rakeysh', )]" 329 | ] 330 | }, 331 | "execution_count": 31, 332 | "metadata": {}, 333 | "output_type": "execute_result" 334 | } 335 | ], 336 | "source": [ 337 | "list(model.vocab.items())[-500:][:100]" 338 | ] 339 | }, 340 | { 341 | "cell_type": "code", 342 | "execution_count": null, 343 | "metadata": {}, 344 | "outputs": [], 345 | "source": [] 346 | } 347 | ], 348 | "metadata": { 349 | "kernelspec": { 350 | "display_name": "Python 3", 351 | "language": "python", 352 | "name": "python3" 353 | }, 354 | "language_info": { 355 | "codemirror_mode": { 356 | "name": "ipython", 357 | "version": 3 358 | }, 359 | "file_extension": ".py", 360 | "mimetype": "text/x-python", 361 | "name": "python", 362 | "nbconvert_exporter": "python", 363 | "pygments_lexer": "ipython3", 364 | "version": "3.6.4" 365 | } 366 | }, 367 | "nbformat": 4, 368 | "nbformat_minor": 2 369 | } 370 | --------------------------------------------------------------------------------