forked from docs/doc-exports
Wuwan, Qi
89983bef36
Reviewed-by: Kacur, Michal <michal.kacur@t-systems.com> Co-authored-by: Wuwan, Qi <wuwanqi1@noreply.gitea.eco.tsi-dev.otc-service.com> Co-committed-by: Wuwan, Qi <wuwanqi1@noreply.gitea.eco.tsi-dev.otc-service.com>
330 lines
27 KiB
HTML
330 lines
27 KiB
HTML
<a name="css_01_0129"></a><a name="css_01_0129"></a>
|
|
|
|
<h1 class="topictitle1">Sample Code for Vector Search on a Client</h1>
|
|
<div id="body0000001261749932"><p id="css_01_0129__en-us_topic_0000001261749932_p143381112172510">Elasticsearch provides standard REST APIs and clients developed using Java, Python, and Go.</p>
|
|
<p id="css_01_0129__en-us_topic_0000001261749932_p1675193372415">Based on the open-source dataset <strong id="css_01_0129__en-us_topic_0000001261749932_b215561919574">SIFT1M</strong> (http://corpus-texmex.irisa.fr/) and Python Elasticsearch client, this section provides a code snippet for creating a vector index, importing vector data, and querying vector data on the client.</p>
|
|
<div class="section" id="css_01_0129__en-us_topic_0000001261749932_section1864223419264"><h4 class="sectiontitle">Prerequisites</h4><p id="css_01_0129__en-us_topic_0000001261749932_p1299195810328">The Python dependency package has been installed on the client. If it is not installed, run the following commands to install it:</p>
|
|
<pre class="screen" id="css_01_0129__en-us_topic_0000001261749932_screen1932333123212">pip install numpy
|
|
pip install elasticsearch==7.6.0</pre>
|
|
</div>
|
|
<div class="section" id="css_01_0129__en-us_topic_0000001261749932_section10885342132619"><h4 class="sectiontitle">Sample Code</h4><div class="codecoloring" codetype="Python" id="css_01_0129__en-us_topic_0000001261749932_screen8396163917303"><div class="highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span class="normal"> 1</span>
|
|
<span class="normal"> 2</span>
|
|
<span class="normal"> 3</span>
|
|
<span class="normal"> 4</span>
|
|
<span class="normal"> 5</span>
|
|
<span class="normal"> 6</span>
|
|
<span class="normal"> 7</span>
|
|
<span class="normal"> 8</span>
|
|
<span class="normal"> 9</span>
|
|
<span class="normal"> 10</span>
|
|
<span class="normal"> 11</span>
|
|
<span class="normal"> 12</span>
|
|
<span class="normal"> 13</span>
|
|
<span class="normal"> 14</span>
|
|
<span class="normal"> 15</span>
|
|
<span class="normal"> 16</span>
|
|
<span class="normal"> 17</span>
|
|
<span class="normal"> 18</span>
|
|
<span class="normal"> 19</span>
|
|
<span class="normal"> 20</span>
|
|
<span class="normal"> 21</span>
|
|
<span class="normal"> 22</span>
|
|
<span class="normal"> 23</span>
|
|
<span class="normal"> 24</span>
|
|
<span class="normal"> 25</span>
|
|
<span class="normal"> 26</span>
|
|
<span class="normal"> 27</span>
|
|
<span class="normal"> 28</span>
|
|
<span class="normal"> 29</span>
|
|
<span class="normal"> 30</span>
|
|
<span class="normal"> 31</span>
|
|
<span class="normal"> 32</span>
|
|
<span class="normal"> 33</span>
|
|
<span class="normal"> 34</span>
|
|
<span class="normal"> 35</span>
|
|
<span class="normal"> 36</span>
|
|
<span class="normal"> 37</span>
|
|
<span class="normal"> 38</span>
|
|
<span class="normal"> 39</span>
|
|
<span class="normal"> 40</span>
|
|
<span class="normal"> 41</span>
|
|
<span class="normal"> 42</span>
|
|
<span class="normal"> 43</span>
|
|
<span class="normal"> 44</span>
|
|
<span class="normal"> 45</span>
|
|
<span class="normal"> 46</span>
|
|
<span class="normal"> 47</span>
|
|
<span class="normal"> 48</span>
|
|
<span class="normal"> 49</span>
|
|
<span class="normal"> 50</span>
|
|
<span class="normal"> 51</span>
|
|
<span class="normal"> 52</span>
|
|
<span class="normal"> 53</span>
|
|
<span class="normal"> 54</span>
|
|
<span class="normal"> 55</span>
|
|
<span class="normal"> 56</span>
|
|
<span class="normal"> 57</span>
|
|
<span class="normal"> 58</span>
|
|
<span class="normal"> 59</span>
|
|
<span class="normal"> 60</span>
|
|
<span class="normal"> 61</span>
|
|
<span class="normal"> 62</span>
|
|
<span class="normal"> 63</span>
|
|
<span class="normal"> 64</span>
|
|
<span class="normal"> 65</span>
|
|
<span class="normal"> 66</span>
|
|
<span class="normal"> 67</span>
|
|
<span class="normal"> 68</span>
|
|
<span class="normal"> 69</span>
|
|
<span class="normal"> 70</span>
|
|
<span class="normal"> 71</span>
|
|
<span class="normal"> 72</span>
|
|
<span class="normal"> 73</span>
|
|
<span class="normal"> 74</span>
|
|
<span class="normal"> 75</span>
|
|
<span class="normal"> 76</span>
|
|
<span class="normal"> 77</span>
|
|
<span class="normal"> 78</span>
|
|
<span class="normal"> 79</span>
|
|
<span class="normal"> 80</span>
|
|
<span class="normal"> 81</span>
|
|
<span class="normal"> 82</span>
|
|
<span class="normal"> 83</span>
|
|
<span class="normal"> 84</span>
|
|
<span class="normal"> 85</span>
|
|
<span class="normal"> 86</span>
|
|
<span class="normal"> 87</span>
|
|
<span class="normal"> 88</span>
|
|
<span class="normal"> 89</span>
|
|
<span class="normal"> 90</span>
|
|
<span class="normal"> 91</span>
|
|
<span class="normal"> 92</span>
|
|
<span class="normal"> 93</span>
|
|
<span class="normal"> 94</span>
|
|
<span class="normal"> 95</span>
|
|
<span class="normal"> 96</span>
|
|
<span class="normal"> 97</span>
|
|
<span class="normal"> 98</span>
|
|
<span class="normal"> 99</span>
|
|
<span class="normal">100</span>
|
|
<span class="normal">101</span>
|
|
<span class="normal">102</span>
|
|
<span class="normal">103</span>
|
|
<span class="normal">104</span>
|
|
<span class="normal">105</span>
|
|
<span class="normal">106</span>
|
|
<span class="normal">107</span>
|
|
<span class="normal">108</span>
|
|
<span class="normal">109</span>
|
|
<span class="normal">110</span>
|
|
<span class="normal">111</span>
|
|
<span class="normal">112</span>
|
|
<span class="normal">113</span>
|
|
<span class="normal">114</span>
|
|
<span class="normal">115</span>
|
|
<span class="normal">116</span>
|
|
<span class="normal">117</span>
|
|
<span class="normal">118</span>
|
|
<span class="normal">119</span>
|
|
<span class="normal">120</span>
|
|
<span class="normal">121</span>
|
|
<span class="normal">122</span>
|
|
<span class="normal">123</span>
|
|
<span class="normal">124</span>
|
|
<span class="normal">125</span>
|
|
<span class="normal">126</span>
|
|
<span class="normal">127</span>
|
|
<span class="normal">128</span>
|
|
<span class="normal">129</span>
|
|
<span class="normal">130</span>
|
|
<span class="normal">131</span>
|
|
<span class="normal">132</span>
|
|
<span class="normal">133</span>
|
|
<span class="normal">134</span>
|
|
<span class="normal">135</span>
|
|
<span class="normal">136</span>
|
|
<span class="normal">137</span>
|
|
<span class="normal">138</span>
|
|
<span class="normal">139</span>
|
|
<span class="normal">140</span>
|
|
<span class="normal">141</span>
|
|
<span class="normal">142</span>
|
|
<span class="normal">143</span>
|
|
<span class="normal">144</span>
|
|
<span class="normal">145</span>
|
|
<span class="normal">146</span>
|
|
<span class="normal">147</span>
|
|
<span class="normal">148</span>
|
|
<span class="normal">149</span>
|
|
<span class="normal">150</span>
|
|
<span class="normal">151</span>
|
|
<span class="normal">152</span>
|
|
<span class="normal">153</span>
|
|
<span class="normal">154</span>
|
|
<span class="normal">155</span></pre></div></td><td class="code"><div><pre><span></span><span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
|
|
<span class="kn">import</span> <span class="nn">time</span>
|
|
<span class="kn">import</span> <span class="nn">json</span>
|
|
|
|
<span class="kn">from</span> <span class="nn">concurrent.futures</span> <span class="kn">import</span> <span class="n">ThreadPoolExecutor</span><span class="p">,</span> <span class="n">wait</span>
|
|
<span class="kn">from</span> <span class="nn">elasticsearch</span> <span class="kn">import</span> <span class="n">Elasticsearch</span>
|
|
<span class="kn">from</span> <span class="nn">elasticsearch</span> <span class="kn">import</span> <span class="n">helpers</span>
|
|
|
|
<span class="n">endpoint</span> <span class="o">=</span> <span class="s1">'http://xxx.xxx.xxx.xxx:9200/'</span>
|
|
|
|
<span class="c1"># Construct an Elasticsearch client object</span>
|
|
<span class="n">es</span> <span class="o">=</span> <span class="n">Elasticsearch</span><span class="p">(</span><span class="n">endpoint</span><span class="p">)</span>
|
|
|
|
<span class="c1"># Index mapping information</span>
|
|
<span class="n">index_mapping</span> <span class="o">=</span> <span class="s1">'''</span>
|
|
<span class="s1">{</span>
|
|
<span class="s1"> "settings": {</span>
|
|
<span class="s1"> "index": {</span>
|
|
<span class="s1"> "vector": "true"</span>
|
|
<span class="s1"> }</span>
|
|
<span class="s1"> },</span>
|
|
<span class="s1"> "mappings": {</span>
|
|
<span class="s1"> "properties": {</span>
|
|
<span class="s1"> "my_vector": {</span>
|
|
<span class="s1"> "type": "vector",</span>
|
|
<span class="s1"> "dimension": 128,</span>
|
|
<span class="s1"> "indexing": true,</span>
|
|
<span class="s1"> "algorithm": "GRAPH",</span>
|
|
<span class="s1"> "metric": "euclidean"</span>
|
|
<span class="s1"> }</span>
|
|
<span class="s1"> }</span>
|
|
<span class="s1"> }</span>
|
|
<span class="s1">}</span>
|
|
<span class="s1">'''</span>
|
|
|
|
<span class="c1"># Create an index.</span>
|
|
<span class="k">def</span> <span class="nf">create_index</span><span class="p">(</span><span class="n">index_name</span><span class="p">,</span> <span class="n">mapping</span><span class="p">):</span>
|
|
<span class="n">res</span> <span class="o">=</span> <span class="n">es</span><span class="o">.</span><span class="n">indices</span><span class="o">.</span><span class="n">create</span><span class="p">(</span><span class="n">index</span><span class="o">=</span><span class="n">index_name</span><span class="p">,</span> <span class="n">ignore</span><span class="o">=</span><span class="mi">400</span><span class="p">,</span> <span class="n">body</span><span class="o">=</span><span class="n">mapping</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">res</span><span class="p">)</span>
|
|
|
|
<span class="c1"># Delete an index.</span>
|
|
<span class="k">def</span> <span class="nf">delete_index</span><span class="p">(</span><span class="n">index_name</span><span class="p">):</span>
|
|
<span class="n">res</span> <span class="o">=</span> <span class="n">es</span><span class="o">.</span><span class="n">indices</span><span class="o">.</span><span class="n">delete</span><span class="p">(</span><span class="n">index</span><span class="o">=</span><span class="n">index_name</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">res</span><span class="p">)</span>
|
|
|
|
|
|
<span class="c1"># Refresh indexes.</span>
|
|
<span class="k">def</span> <span class="nf">refresh_index</span><span class="p">(</span><span class="n">index_name</span><span class="p">):</span>
|
|
<span class="n">res</span> <span class="o">=</span> <span class="n">es</span><span class="o">.</span><span class="n">indices</span><span class="o">.</span><span class="n">refresh</span><span class="p">(</span><span class="n">index</span><span class="o">=</span><span class="n">index_name</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="n">res</span><span class="p">)</span>
|
|
|
|
|
|
<span class="c1"># Merge index segments.</span>
|
|
<span class="k">def</span> <span class="nf">merge_index</span><span class="p">(</span><span class="n">index_name</span><span class="p">,</span> <span class="n">seg_cnt</span><span class="o">=</span><span class="mi">1</span><span class="p">):</span>
|
|
<span class="n">start</span> <span class="o">=</span> <span class="n">time</span><span class="o">.</span><span class="n">time</span><span class="p">()</span>
|
|
<span class="n">es</span><span class="o">.</span><span class="n">indices</span><span class="o">.</span><span class="n">forcemerge</span><span class="p">(</span><span class="n">index</span><span class="o">=</span><span class="n">index_name</span><span class="p">,</span> <span class="n">max_num_segments</span><span class="o">=</span><span class="n">seg_cnt</span><span class="p">,</span> <span class="n">request_timeout</span><span class="o">=</span><span class="mi">36000</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">" Complete the merge within </span><span class="si">{</span><span class="n">time</span><span class="o">.</span><span class="n">time</span><span class="p">()</span><span class="w"> </span><span class="o">-</span><span class="w"> </span><span class="n">start</span><span class="si">}</span><span class="s2"> seconds"</span><span class="p">)</span>
|
|
|
|
|
|
<span class="c1"># Load vector data.</span>
|
|
<span class="k">def</span> <span class="nf">load_vectors</span><span class="p">(</span><span class="n">file_name</span><span class="p">):</span>
|
|
<span class="n">fv</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">fromfile</span><span class="p">(</span><span class="n">file_name</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>
|
|
<span class="n">dim</span> <span class="o">=</span> <span class="n">fv</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">int32</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
|
|
<span class="n">vectors</span> <span class="o">=</span> <span class="n">fv</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span> <span class="o">+</span> <span class="n">dim</span><span class="p">)[:,</span> <span class="mi">1</span><span class="p">:]</span>
|
|
<span class="k">return</span> <span class="n">vectors</span>
|
|
|
|
|
|
<span class="c1"># Load the ground_truth data.</span>
|
|
<span class="k">def</span> <span class="nf">load_gts</span><span class="p">(</span><span class="n">file_name</span><span class="p">):</span>
|
|
<span class="n">fv</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">fromfile</span><span class="p">(</span><span class="n">file_name</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">np</span><span class="o">.</span><span class="n">int32</span><span class="p">)</span>
|
|
<span class="n">dim</span> <span class="o">=</span> <span class="n">fv</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">int32</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span>
|
|
<span class="n">gts</span> <span class="o">=</span> <span class="n">fv</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span> <span class="o">+</span> <span class="n">dim</span><span class="p">)[:,</span> <span class="mi">1</span><span class="p">:]</span>
|
|
<span class="k">return</span> <span class="n">gts</span>
|
|
|
|
|
|
<span class="k">def</span> <span class="nf">partition</span><span class="p">(</span><span class="n">ls</span><span class="p">,</span> <span class="n">size</span><span class="p">):</span>
|
|
<span class="k">return</span> <span class="p">[</span><span class="n">ls</span><span class="p">[</span><span class="n">i</span><span class="p">:</span><span class="n">i</span> <span class="o">+</span> <span class="n">size</span><span class="p">]</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="nb">len</span><span class="p">(</span><span class="n">ls</span><span class="p">),</span> <span class="n">size</span><span class="p">)]</span>
|
|
|
|
|
|
<span class="c1"># Write vector data.</span>
|
|
<span class="k">def</span> <span class="nf">write_index</span><span class="p">(</span><span class="n">index_name</span><span class="p">,</span> <span class="n">vec_file</span><span class="p">):</span>
|
|
<span class="n">pool</span> <span class="o">=</span> <span class="n">ThreadPoolExecutor</span><span class="p">(</span><span class="n">max_workers</span><span class="o">=</span><span class="mi">8</span><span class="p">)</span>
|
|
<span class="n">tasks</span> <span class="o">=</span> <span class="p">[]</span>
|
|
|
|
<span class="n">vectors</span> <span class="o">=</span> <span class="n">load_vectors</span><span class="p">(</span><span class="n">vec_file</span><span class="p">)</span>
|
|
<span class="n">bulk_size</span> <span class="o">=</span> <span class="mi">1000</span>
|
|
<span class="n">partitions</span> <span class="o">=</span> <span class="n">partition</span><span class="p">(</span><span class="n">vectors</span><span class="p">,</span> <span class="n">bulk_size</span><span class="p">)</span>
|
|
|
|
<span class="n">start</span> <span class="o">=</span> <span class="n">time</span><span class="o">.</span><span class="n">time</span><span class="p">()</span>
|
|
<span class="n">start_id</span> <span class="o">=</span> <span class="mi">0</span>
|
|
<span class="k">for</span> <span class="n">vecs</span> <span class="ow">in</span> <span class="n">partitions</span><span class="p">:</span>
|
|
<span class="n">tasks</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">pool</span><span class="o">.</span><span class="n">submit</span><span class="p">(</span><span class="n">write_bulk</span><span class="p">,</span> <span class="n">index_name</span><span class="p">,</span> <span class="n">vecs</span><span class="p">,</span> <span class="n">start_id</span><span class="p">))</span>
|
|
<span class="n">start_id</span> <span class="o">+=</span> <span class="nb">len</span><span class="p">(</span><span class="n">vecs</span><span class="p">)</span>
|
|
<span class="n">wait</span><span class="p">(</span><span class="n">tasks</span><span class="p">)</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">" Complete the writing within </span><span class="si">{</span><span class="n">time</span><span class="o">.</span><span class="n">time</span><span class="p">()</span><span class="w"> </span><span class="o">-</span><span class="w"> </span><span class="n">start</span><span class="si">}</span><span class="s2"> seconds"</span><span class="p">)</span>
|
|
|
|
|
|
<span class="k">def</span> <span class="nf">write_bulk</span><span class="p">(</span><span class="n">index_name</span><span class="p">,</span> <span class="n">vecs</span><span class="p">,</span> <span class="n">start_id</span><span class="p">):</span>
|
|
<span class="n">actions</span> <span class="o">=</span> <span class="p">[</span>
|
|
<span class="p">{</span>
|
|
<span class="s2">"_index"</span><span class="p">:</span> <span class="n">index_name</span><span class="p">,</span>
|
|
<span class="s2">"my_vector"</span><span class="p">:</span> <span class="n">vecs</span><span class="p">[</span><span class="n">j</span><span class="p">]</span><span class="o">.</span><span class="n">tolist</span><span class="p">(),</span>
|
|
<span class="s2">"_id"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">j</span> <span class="o">+</span> <span class="n">start_id</span><span class="p">)</span>
|
|
<span class="p">}</span>
|
|
<span class="k">for</span> <span class="n">j</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">vecs</span><span class="p">))</span>
|
|
<span class="p">]</span>
|
|
<span class="n">helpers</span><span class="o">.</span><span class="n">bulk</span><span class="p">(</span><span class="n">es</span><span class="p">,</span> <span class="n">actions</span><span class="p">,</span> <span class="n">request_timeout</span><span class="o">=</span><span class="mi">3600</span><span class="p">)</span>
|
|
|
|
|
|
<span class="c1"># Query an index.</span>
|
|
<span class="k">def</span> <span class="nf">search_index</span><span class="p">(</span><span class="n">index_name</span><span class="p">,</span> <span class="n">query_file</span><span class="p">,</span> <span class="n">gt_file</span><span class="p">,</span> <span class="n">k</span><span class="p">):</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="s2">"Start query! Index name: "</span> <span class="o">+</span> <span class="n">index_name</span><span class="p">)</span>
|
|
|
|
<span class="n">queries</span> <span class="o">=</span> <span class="n">load_vectors</span><span class="p">(</span><span class="n">query_file</span><span class="p">)</span>
|
|
<span class="n">gt</span> <span class="o">=</span> <span class="n">load_gts</span><span class="p">(</span><span class="n">gt_file</span><span class="p">)</span>
|
|
|
|
<span class="n">took</span> <span class="o">=</span> <span class="mi">0</span>
|
|
<span class="n">precision</span> <span class="o">=</span> <span class="p">[]</span>
|
|
<span class="k">for</span> <span class="n">idx</span><span class="p">,</span> <span class="n">query</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">queries</span><span class="p">):</span>
|
|
<span class="n">hits</span> <span class="o">=</span> <span class="nb">set</span><span class="p">()</span>
|
|
<span class="n">query_json</span> <span class="o">=</span> <span class="p">{</span>
|
|
<span class="s2">"size"</span><span class="p">:</span> <span class="n">k</span><span class="p">,</span>
|
|
<span class="s2">"_source"</span><span class="p">:</span> <span class="kc">False</span><span class="p">,</span>
|
|
<span class="s2">"query"</span><span class="p">:</span> <span class="p">{</span>
|
|
<span class="s2">"vector"</span><span class="p">:</span> <span class="p">{</span>
|
|
<span class="s2">"my_vector"</span><span class="p">:</span> <span class="p">{</span>
|
|
<span class="s2">"vector"</span><span class="p">:</span> <span class="n">query</span><span class="o">.</span><span class="n">tolist</span><span class="p">(),</span>
|
|
<span class="s2">"topk"</span><span class="p">:</span> <span class="n">k</span>
|
|
<span class="p">}</span>
|
|
<span class="p">}</span>
|
|
<span class="p">}</span>
|
|
<span class="p">}</span>
|
|
<span class="n">res</span> <span class="o">=</span> <span class="n">es</span><span class="o">.</span><span class="n">search</span><span class="p">(</span><span class="n">index</span><span class="o">=</span><span class="n">index_name</span><span class="p">,</span> <span class="n">body</span><span class="o">=</span><span class="n">json</span><span class="o">.</span><span class="n">dumps</span><span class="p">(</span><span class="n">query_json</span><span class="p">))</span>
|
|
|
|
<span class="k">for</span> <span class="n">hit</span> <span class="ow">in</span> <span class="n">res</span><span class="p">[</span><span class="s1">'hits'</span><span class="p">][</span><span class="s1">'hits'</span><span class="p">]:</span>
|
|
<span class="n">hits</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="nb">int</span><span class="p">(</span><span class="n">hit</span><span class="p">[</span><span class="s1">'_id'</span><span class="p">]))</span>
|
|
<span class="n">precision</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">hits</span><span class="o">.</span><span class="n">intersection</span><span class="p">(</span><span class="nb">set</span><span class="p">(</span><span class="n">gt</span><span class="p">[</span><span class="n">idx</span><span class="p">,</span> <span class="p">:</span><span class="n">k</span><span class="p">])))</span> <span class="o">/</span> <span class="n">k</span><span class="p">)</span>
|
|
<span class="n">took</span> <span class="o">+=</span> <span class="n">res</span><span class="p">[</span><span class="s1">'took'</span><span class="p">]</span>
|
|
|
|
<span class="nb">print</span><span class="p">(</span><span class="s2">"precision: "</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="nb">sum</span><span class="p">(</span><span class="n">precision</span><span class="p">)</span> <span class="o">/</span> <span class="nb">len</span><span class="p">(</span><span class="n">precision</span><span class="p">)))</span>
|
|
<span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">" Complete the retrieval within </span><span class="si">{</span><span class="n">took</span><span class="w"> </span><span class="o">/</span><span class="w"> </span><span class="mi">1000</span><span class="si">:</span><span class="s2">.2f</span><span class="si">}</span><span class="s2"> seconds; average took size is </span><span class="si">{</span><span class="n">took</span><span class="w"> </span><span class="o">/</span><span class="w"> </span><span class="nb">len</span><span class="p">(</span><span class="n">queries</span><span class="p">)</span><span class="si">:</span><span class="s2">.2f</span><span class="si">}</span><span class="s2"> ms"</span><span class="p">)</span>
|
|
|
|
|
|
<span class="k">if</span> <span class="vm">__name__</span> <span class="o">==</span> <span class="s2">"__main__"</span><span class="p">:</span>
|
|
<span class="n">vec_file</span> <span class="o">=</span> <span class="sa">r</span><span class="s2">"./data/sift/sift_base.fvecs"</span>
|
|
<span class="n">qry_file</span> <span class="o">=</span> <span class="sa">r</span><span class="s2">"./data/sift/sift_query.fvecs"</span>
|
|
<span class="n">gt_file</span> <span class="o">=</span> <span class="sa">r</span><span class="s2">"./data/sift/sift_groundtruth.ivecs"</span>
|
|
|
|
<span class="n">index</span> <span class="o">=</span> <span class="s2">"test"</span>
|
|
<span class="n">create_index</span><span class="p">(</span><span class="n">index</span><span class="p">,</span> <span class="n">index_mapping</span><span class="p">)</span>
|
|
<span class="n">write_index</span><span class="p">(</span><span class="n">index</span><span class="p">,</span> <span class="n">vec_file</span><span class="p">)</span>
|
|
<span class="n">merge_index</span><span class="p">(</span><span class="n">index</span><span class="p">)</span>
|
|
<span class="n">refresh_index</span><span class="p">(</span><span class="n">index</span><span class="p">)</span>
|
|
|
|
<span class="n">search_index</span><span class="p">(</span><span class="n">index</span><span class="p">,</span> <span class="n">qry_file</span><span class="p">,</span> <span class="n">gt_file</span><span class="p">,</span> <span class="mi">10</span><span class="p">)</span>
|
|
</pre></div></td></tr></table></div>
|
|
|
|
</div>
|
|
</div>
|
|
</div>
|
|
<div>
|
|
<div class="familylinks">
|
|
<div class="parentlink"><strong>Parent topic:</strong> <a href="css_01_0117.html">Vector Retrieval</a></div>
|
|
</div>
|
|
</div>
|
|
|