AlCatt91 published a site update

graphcore-research · Jan 8, 2024 · b28953d · b28953d
1 parent 96339f1
commit b28953d
Show file tree

Hide file tree

Showing 6 changed files with 52 additions and 13 deletions.
diff --git a/_modules/besskge/dataset.html b/_modules/besskge/dataset.html
@@ -110,6 +110,11 @@ <h1>Source code for besskge.dataset</h1><div class="highlight"><pre>
     <span class="c1">#: {part: int32[n_triple, {h,r,t}]}</span>
     <span class="n">triples</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">NDArray</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="c1">#: IDs of the triples in KGDataset.triples wrt</span>
+    <span class="c1">#: the ordering in the original array/dataframe</span>
+    <span class="c1">#: from where the triples originate.</span>
+    <span class="n">original_triple_ids</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">NDArray</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="c1">#: Entity labels by ID; str[n_entity]</span>
     <span class="n">entity_dict</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span>
 
@@ -164,6 +169,9 @@ <h1>Source code for besskge.dataset</h1><div class="highlight"><pre>
 <span class="sd">        and relations have already been assigned. Note that, if entities have</span>
 <span class="sd">        types, entities of the same type need to have contiguous IDs.</span>
 <span class="sd">        Triples are randomly split in train/validation/test sets.</span>
+<span class="sd">        The attribute `KGDataset.original_triple_ids` stores the IDs</span>
+<span class="sd">        of the triples in each split wrt the original ordering in `data`.</span>
+
 <span class="sd">        If a pre-defined train/validation/test split is wanted, the KGDataset</span>
 <span class="sd">        class should be instantiated manually.</span>
 
@@ -188,21 +196,27 @@ <h1>Source code for besskge.dataset</h1><div class="highlight"><pre>
         <span class="n">num_valid</span> <span class="o">=</span> <span class="nb">int</span><span class="p">(</span><span class="n">num_triples</span> <span class="o">*</span> <span class="n">split</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
 
         <span class="n">rng</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">default_rng</span><span class="p">(</span><span class="n">seed</span><span class="o">=</span><span class="n">seed</span><span class="p">)</span>
-        <span class="n">rng</span><span class="o">.</span><span class="n">shuffle</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
-
-        <span class="n">triples</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
-        <span class="n">triples</span><span class="p">[</span><span class="s2">&quot;train&quot;</span><span class="p">],</span> <span class="n">triples</span><span class="p">[</span><span class="s2">&quot;valid&quot;</span><span class="p">],</span> <span class="n">triples</span><span class="p">[</span><span class="s2">&quot;test&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">split</span><span class="p">(</span>
-            <span class="n">data</span><span class="p">,</span> <span class="p">(</span><span class="n">num_train</span><span class="p">,</span> <span class="n">num_train</span> <span class="o">+</span> <span class="n">num_valid</span><span class="p">),</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</span>
+        <span class="n">id_shuffle</span> <span class="o">=</span> <span class="n">rng</span><span class="o">.</span><span class="n">permutation</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">num_triples</span><span class="p">))</span>
+        <span class="n">triple_ids</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
+        <span class="n">triple_ids</span><span class="p">[</span><span class="s2">&quot;train&quot;</span><span class="p">],</span> <span class="n">triple_ids</span><span class="p">[</span><span class="s2">&quot;valid&quot;</span><span class="p">],</span> <span class="n">triple_ids</span><span class="p">[</span><span class="s2">&quot;test&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">split</span><span class="p">(</span>
+            <span class="n">id_shuffle</span><span class="p">,</span> <span class="p">(</span><span class="n">num_train</span><span class="p">,</span> <span class="n">num_train</span> <span class="o">+</span> <span class="n">num_valid</span><span class="p">),</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</span>
         <span class="p">)</span>
+        <span class="n">triples</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span>
+        <span class="n">triples</span><span class="p">[</span><span class="s2">&quot;train&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">data</span><span class="p">[</span><span class="n">triple_ids</span><span class="p">[</span><span class="s2">&quot;train&quot;</span><span class="p">]]</span>
+        <span class="n">triples</span><span class="p">[</span><span class="s2">&quot;valid&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">data</span><span class="p">[</span><span class="n">triple_ids</span><span class="p">[</span><span class="s2">&quot;valid&quot;</span><span class="p">]]</span>
+        <span class="n">triples</span><span class="p">[</span><span class="s2">&quot;test&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">data</span><span class="p">[</span><span class="n">triple_ids</span><span class="p">[</span><span class="s2">&quot;test&quot;</span><span class="p">]]</span>
 
-        <span class="k">return</span> <span class="bp">cls</span><span class="p">(</span>
+        <span class="n">ds</span> <span class="o">=</span> <span class="bp">cls</span><span class="p">(</span>
             <span class="n">n_entity</span><span class="o">=</span><span class="n">data</span><span class="p">[:,</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">]]</span><span class="o">.</span><span class="n">max</span><span class="p">()</span> <span class="o">+</span> <span class="mi">1</span><span class="p">,</span>
             <span class="n">n_relation_type</span><span class="o">=</span><span class="n">data</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">max</span><span class="p">()</span> <span class="o">+</span> <span class="mi">1</span><span class="p">,</span>
             <span class="n">entity_dict</span><span class="o">=</span><span class="n">entity_dict</span><span class="p">,</span>
             <span class="n">relation_dict</span><span class="o">=</span><span class="n">relation_dict</span><span class="p">,</span>
             <span class="n">type_offsets</span><span class="o">=</span><span class="n">type_offsets</span><span class="p">,</span>
             <span class="n">triples</span><span class="o">=</span><span class="n">triples</span><span class="p">,</span>
-        <span class="p">)</span></div>
+            <span class="n">original_triple_ids</span><span class="o">=</span><span class="n">triple_ids</span><span class="p">,</span>
+        <span class="p">)</span>
+
+        <span class="k">return</span> <span class="n">ds</span></div>
 
 <div class="viewcode-block" id="KGDataset.from_dataframe"><a class="viewcode-back" href="../../generated/besskge.dataset.KGDataset.html#besskge.dataset.KGDataset.from_dataframe">[docs]</a>    <span class="nd">@classmethod</span>
     <span class="k">def</span> <span class="nf">from_dataframe</span><span class="p">(</span>
@@ -293,6 +307,9 @@ <h1>Source code for besskge.dataset</h1><div class="highlight"><pre>
                 <span class="n">relation_dict</span><span class="o">=</span><span class="n">relation_dict</span><span class="p">,</span>
                 <span class="n">type_offsets</span><span class="o">=</span><span class="n">type_offsets</span><span class="p">,</span>
                 <span class="n">triples</span><span class="o">=</span><span class="n">triples</span><span class="p">,</span>
+                <span class="n">original_triple_ids</span><span class="o">=</span><span class="p">{</span>
+                    <span class="n">k</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">v</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="k">for</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span> <span class="ow">in</span> <span class="n">triples</span><span class="o">.</span><span class="n">items</span><span class="p">()</span>
+                <span class="p">},</span>
             <span class="p">)</span></div>
 
 <div class="viewcode-block" id="KGDataset.build_ogbl_biokg"><a class="viewcode-back" href="../../generated/besskge.dataset.KGDataset.html#besskge.dataset.KGDataset.build_ogbl_biokg">[docs]</a>    <span class="nd">@classmethod</span>
@@ -354,6 +371,7 @@ <h1>Source code for besskge.dataset</h1><div class="highlight"><pre>
             <span class="n">relation_dict</span><span class="o">=</span><span class="n">rel_dict</span><span class="p">,</span>
             <span class="n">type_offsets</span><span class="o">=</span><span class="n">type_offsets</span><span class="p">,</span>
             <span class="n">triples</span><span class="o">=</span><span class="n">triples</span><span class="p">,</span>
+            <span class="n">original_triple_ids</span><span class="o">=</span><span class="p">{</span><span class="n">k</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">v</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="k">for</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span> <span class="ow">in</span> <span class="n">triples</span><span class="o">.</span><span class="n">items</span><span class="p">()},</span>
             <span class="n">neg_heads</span><span class="o">=</span><span class="n">neg_heads</span><span class="p">,</span>
             <span class="n">neg_tails</span><span class="o">=</span><span class="n">neg_tails</span><span class="p">,</span>
         <span class="p">)</span></div>
@@ -403,6 +421,7 @@ <h1>Source code for besskge.dataset</h1><div class="highlight"><pre>
             <span class="n">relation_dict</span><span class="o">=</span><span class="n">rel_dict</span><span class="p">,</span>
             <span class="n">type_offsets</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
             <span class="n">triples</span><span class="o">=</span><span class="n">triples</span><span class="p">,</span>
+            <span class="n">original_triple_ids</span><span class="o">=</span><span class="p">{</span><span class="n">k</span><span class="p">:</span> <span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">v</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="k">for</span> <span class="n">k</span><span class="p">,</span> <span class="n">v</span> <span class="ow">in</span> <span class="n">triples</span><span class="o">.</span><span class="n">items</span><span class="p">()},</span>
             <span class="n">neg_heads</span><span class="o">=</span><span class="n">neg_heads</span><span class="p">,</span>
             <span class="n">neg_tails</span><span class="o">=</span><span class="n">neg_tails</span><span class="p">,</span>
         <span class="p">)</span></div>