diff --git a/examples/04_Drug_Names_ATC.ipynb b/examples/04_Drug_Names_ATC.ipynb
index 1c3db21..48a4df0 100644
--- a/examples/04_Drug_Names_ATC.ipynb
+++ b/examples/04_Drug_Names_ATC.ipynb
@@ -56,11 +56,11 @@
{
"data": {
"text/html": [
- "[03/07/24 10:43:03] INFO Loading hierarchical faiss index sap_bert_linker.py:153\n",
+ "[03/07/24 18:18:49] INFO Loading hierarchical faiss index sap_bert_linker.py:153\n",
"\n"
],
"text/plain": [
- "\u001b[2;36m[03/07/24 10:43:03]\u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO \u001b[0m Loading hierarchical faiss index \u001b]8;id=280662;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/sap_bert_linker.py\u001b\\\u001b[2msap_bert_linker.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=664469;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/sap_bert_linker.py#153\u001b\\\u001b[2m153\u001b[0m\u001b]8;;\u001b\\\n"
+ "\u001b[2;36m[03/07/24 18:18:49]\u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO \u001b[0m Loading hierarchical faiss index \u001b]8;id=789095;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/sap_bert_linker.py\u001b\\\u001b[2msap_bert_linker.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=471624;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/sap_bert_linker.py#153\u001b\\\u001b[2m153\u001b[0m\u001b]8;;\u001b\\\n"
]
},
"metadata": {},
@@ -75,7 +75,7 @@
"
\n"
],
"text/plain": [
- "\u001b[2;36m \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO \u001b[0m Loading index from \u001b]8;id=926064;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/faiss_indexer.py\u001b\\\u001b[2mfaiss_indexer.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=529795;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/faiss_indexer.py#64\u001b\\\u001b[2m64\u001b[0m\u001b]8;;\u001b\\\n",
+ "\u001b[2;36m \u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO \u001b[0m Loading index from \u001b]8;id=852479;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/faiss_indexer.py\u001b\\\u001b[2mfaiss_indexer.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=39591;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/faiss_indexer.py#64\u001b\\\u001b[2m64\u001b[0m\u001b]8;;\u001b\\\n",
"\u001b[2;36m \u001b[0m \u001b[35m/home/Florian.Borchert/.cache/xmen/atc/index/sapbert/\u001b[0m\u001b[95membed_faiss_h\u001b[0m \u001b[2m \u001b[0m\n",
"\u001b[2;36m \u001b[0m \u001b[95mier.pickle\u001b[0m \u001b[2m \u001b[0m\n"
]
@@ -86,12 +86,12 @@
{
"data": {
"text/html": [
- "[03/07/24 10:43:04] INFO Loaded index of type <class 'faiss.swigfaiss.IndexHNSWFlat'> and faiss_indexer.py:66\n",
+ "[03/07/24 18:18:50] INFO Loaded index of type <class 'faiss.swigfaiss.IndexHNSWFlat'> and faiss_indexer.py:66\n",
" size 470941 \n",
"\n"
],
"text/plain": [
- "\u001b[2;36m[03/07/24 10:43:04]\u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO \u001b[0m Loaded index of type \u001b[1m<\u001b[0m\u001b[1;95mclass\u001b[0m\u001b[39m \u001b[0m\u001b[32m'faiss.swigfaiss.IndexHNSWFlat'\u001b[0m\u001b[1m>\u001b[0m and \u001b]8;id=99502;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/faiss_indexer.py\u001b\\\u001b[2mfaiss_indexer.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=2245;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/faiss_indexer.py#66\u001b\\\u001b[2m66\u001b[0m\u001b]8;;\u001b\\\n",
+ "\u001b[2;36m[03/07/24 18:18:50]\u001b[0m\u001b[2;36m \u001b[0m\u001b[34mINFO \u001b[0m Loaded index of type \u001b[1m<\u001b[0m\u001b[1;95mclass\u001b[0m\u001b[39m \u001b[0m\u001b[32m'faiss.swigfaiss.IndexHNSWFlat'\u001b[0m\u001b[1m>\u001b[0m and \u001b]8;id=788991;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/faiss_indexer.py\u001b\\\u001b[2mfaiss_indexer.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=743737;file:///mnt/nfs/home/Florian.Borchert/workspace/xmen/xmen/linkers/faiss_indexer.py#66\u001b\\\u001b[2m66\u001b[0m\u001b]8;;\u001b\\\n",
"\u001b[2;36m \u001b[0m size \u001b[1;36m470941\u001b[0m \u001b[2m \u001b[0m\n"
]
},
@@ -134,7 +134,7 @@
{
"data": {
"application/vnd.jupyter.widget-view+json": {
- "model_id": "cd11b10d37194fa2b3f37d2777bdcb4d",
+ "model_id": "5a4c8dd8146148aaa5694a517d6aa965",
"version_major": 2,
"version_minor": 0
},
@@ -224,7 +224,7 @@
"source": [
"for d, p in zip(drug_mentions, predictions):\n",
" print('Input:', d)\n",
- " top_candidate = p[0]['normalized'][0]\n",
+ " top_candidate = p['normalized'][0]\n",
" print('Confidence:', top_candidate['score'])\n",
" print(kb.cui_to_entity[top_candidate['db_id']])\n",
" print('------')"
diff --git a/xmen/linkers/__init__.py b/xmen/linkers/__init__.py
index 74e0255..1fba005 100644
--- a/xmen/linkers/__init__.py
+++ b/xmen/linkers/__init__.py
@@ -63,10 +63,14 @@ def predict_no_context(
ds = from_spans(entities=spans, sentences=sentences)
result = self.predict_batch(ds, batch_size)
if is_str:
- assert len(result["entities"]) == 1
- return result["entities"][0]
+ assert len(result["entities"]) == 1 and len(result["entities"][0]) == 1
+ return result["entities"][0][0]
else:
- return result["entities"]
+ _result = []
+ for r in result["entities"]:
+ assert len(r) == 1
+ _result.append(r[0])
+ return _result
class RerankedLinker(EntityLinker):