yujunjun/haystack
1
0
Fork 0
mirror of https://github.com/deepset-ai/haystack.git synced 2025-11-01 10:19:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add filtering by tags for InMemoryDocumentStore (#108)

Browse Source
This commit is contained in:
Stan Kirdey 2020-05-14 13:12:25 -07:00 committed by GitHub
parent 73aed42f14
commit 72a3b70d7a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 117 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
README.rst
View File

@ -223,3 +223,8 @@ You will find the Swagger API documentation at http://127.0.0.1:80/docs
* Coming soon: more file formats for document upload, metrics for label quality ...
.. image:: https://raw.githubusercontent.com/deepset-ai/haystack/master/docs/img/annotation_tool.png
7. Development
-------------------
* Unit tests are executed by running ```tox```

7
haystack/database/base.py
View File

@ -38,9 +38,4 @@ class Document(BaseModel):
question: Optional[str] = Field(None, description="Question text for FAQs.")
query_score: Optional[float] = Field(None, description="Elasticsearch query score for a retrieved document")
meta: Optional[Dict[str, Any]] = Field(None, description="")
def __getitem__(self, item):
if item == 'text':
return self.text
if item == 'id':
return self.id
tags: Optional[Dict[str, Any]] = Field(None, description="Tags that allow filtering of the data")

45
haystack/database/memory.py
View File

@ -12,6 +12,10 @@ class InMemoryDocumentStore(BaseDocumentStore):
def write_documents(self, documents):
import hashlib
if documents is None:
return
for document in documents:
name = document.get("name", None)
text = document.get("text", None)
@ -20,10 +24,30 @@ class InMemoryDocumentStore(BaseDocumentStore):
continue
signature = name + text
hash = hashlib.md5(signature.encode("utf-8")).hexdigest()
self.docs[hash] = document
tags = document.get('tags', [])
self._map_tags_to_ids(hash, tags)
def _map_tags_to_ids(self, hash, tags):
if isinstance(tags, list):
for tag in tags:
if isinstance(tag, dict):
tag_keys = tag.keys()
for tag_key in tag_keys:
tag_values = tag.get(tag_key, [])
if tag_values:
for tag_value in tag_values:
comp_key = str((tag_key, tag_value))
if comp_key in self.doc_tags:
self.doc_tags[comp_key].append(hash)
else:
self.doc_tags[comp_key] = [hash]
def get_document_by_id(self, id):
return self.docs[id]
@ -32,10 +56,27 @@ class InMemoryDocumentStore(BaseDocumentStore):
The format for the dict is {"tag-1": "value-1", "tag-2": "value-2" ...}
The format for the dict is {"tag-1": ["value-1","value-2"], "tag-2": ["value-3]" ...}
"""
pass
if not isinstance(tags, list):
tags = [tags]
result = self._find_ids_by_tags(tags)
return result
def _find_ids_by_tags(self, tags):
result = []
for tag in tags:
tag_keys = tag.keys()
for tag_key in tag_keys:
tag_values = tag.get(tag_key, None)
if tag_values:
for tag_value in tag_values:
comp_key = str((tag_key, tag_value))
doc_ids = self.doc_tags.get(comp_key, [])
for doc_id in doc_ids:
result.append(self.docs.get(doc_id))
return result
def get_document_count(self):
return len(self.docs.items())
def get_all_documents(self):
return [Document(id=item[0], text=item[1]['text'], name=item[1]['name']) for item in self.docs.items()]
return [Document(id=item[0], text=item[1]['text'], name=item[1]['name'], meta=item[1].get('meta', {})) for item in self.docs.items()]

1
test/test_db.py
View File

@ -20,6 +20,7 @@ def test_elasticsearch_write_read(elasticsearch_fixture):
write_documents_to_db(document_store=document_store, document_dir="samples/docs")
sleep(2) # wait for documents to be available for query
documents = document_store.get_all_documents()
print(documents)
assert len(documents) == 2
assert documents[0].id
assert documents[0].text

1
test/test_document.py
View File

@ -4,4 +4,3 @@ from haystack.database.base import Document
def test_document_data_access():
doc = Document(id=1, text="test")
assert doc.text == "test"
assert doc['text'] == "test"

2
test/test_farm_reader.py
View File

@ -1,5 +1,3 @@
import pytest
from haystack.reader.farm import FARMReader

70
test/test_in_memory_store.py
View File

@ -5,14 +5,15 @@ from haystack.retriever.tfidf import TfidfRetriever
def test_finder_get_answers_with_in_memory_store():
test_docs = [
{"name": "testing the finder 1", "text": "testing the finder with pyhton unit test 1"},
{"name": "testing the finder 2", "text": "testing the finder with pyhton unit test 2"},
{"name": "testing the finder 3", "text": "testing the finder with pyhton unit test 3"}
{"name": "testing the finder 1", "text": "testing the finder with pyhton unit test 1", 'meta': {'url': 'url'}},
{"name": "testing the finder 2", "text": "testing the finder with pyhton unit test 2", 'meta': {'url': 'url'}},
{"name": "testing the finder 3", "text": "testing the finder with pyhton unit test 3", 'meta': {'url': 'url'}}
]
from haystack.database.memory import InMemoryDocumentStore
document_store = InMemoryDocumentStore()
document_store.write_documents(test_docs)
retriever = TfidfRetriever(document_store=document_store)
reader = TransformersReader(model="distilbert-base-uncased-distilled-squad",
tokenizer="distilbert-base-uncased", use_gpu=-1)
@ -20,3 +21,66 @@ def test_finder_get_answers_with_in_memory_store():
prediction = finder.get_answers(question="testing finder", top_k_retriever=10,
top_k_reader=5)
assert prediction is not None
def test_memory_store_get_by_tags():
test_docs = [
{"name": "testing the finder 1", "text": "testing the finder with pyhton unit test 1", 'meta': {'url': 'url'}},
{"name": "testing the finder 2", "text": "testing the finder with pyhton unit test 2", 'meta': {'url': None}},
{"name": "testing the finder 3", "text": "testing the finder with pyhton unit test 3", 'meta': {'url': 'url'}}
]
from haystack.database.memory import InMemoryDocumentStore
document_store = InMemoryDocumentStore()
document_store.write_documents(test_docs)
docs = document_store.get_document_ids_by_tags({'has_url': 'false'})
assert docs == []
def test_memory_store_get_by_tag_lists_union():
test_docs = [
{"name": "testing the finder 1", "text": "testing the finder with pyhton unit test 1", 'meta': {'url': 'url'}, 'tags': [{'tag2': ["1"]}]},
{"name": "testing the finder 2", "text": "testing the finder with pyhton unit test 2", 'meta': {'url': None}, 'tags': [{'tag1': ['1']}]},
{"name": "testing the finder 3", "text": "testing the finder with pyhton unit test 3", 'meta': {'url': 'url'}, 'tags': [{'tag2': ["1", "2"]}]}
]
from haystack.database.memory import InMemoryDocumentStore
document_store = InMemoryDocumentStore()
document_store.write_documents(test_docs)
docs = document_store.get_document_ids_by_tags({'tag2': ["1"]})
assert docs == [
{'name': 'testing the finder 1', 'text': 'testing the finder with pyhton unit test 1', 'meta': {'url': 'url'}, 'tags': [{'tag2': ['1']}]},
{'name': 'testing the finder 3', 'text': 'testing the finder with pyhton unit test 3', 'meta': {'url': 'url'}, 'tags': [{'tag2': ['1', '2']}]}
]
def test_memory_store_get_by_tag_lists_non_existent_tag():
test_docs = [
{"name": "testing the finder 1", "text": "testing the finder with pyhton unit test 1", 'meta': {'url': 'url'}, 'tags': [{'tag1': ["1"]}]},
]
from haystack.database.memory import InMemoryDocumentStore
document_store = InMemoryDocumentStore()
document_store.write_documents(test_docs)
docs = document_store.get_document_ids_by_tags({'tag1': ["3"]})
assert docs == []
def test_memory_store_get_by_tag_lists_disjoint():
test_docs = [
{"name": "testing the finder 1", "text": "testing the finder with pyhton unit test 1", 'meta': {'url': 'url'}, 'tags': [{'tag1': ["1"]}]},
{"name": "testing the finder 2", "text": "testing the finder with pyhton unit test 2", 'meta': {'url': None}, 'tags': [{'tag2': ['1']}]},
{"name": "testing the finder 3", "text": "testing the finder with pyhton unit test 3", 'meta': {'url': 'url'}, 'tags': [{'tag3': ["1", "2"]}]},
{"name": "testing the finder 4", "text": "testing the finder with pyhton unit test 3", 'meta': {'url': 'url'}, 'tags': [{'tag3': ["1", "3"]}]}
]
from haystack.database.memory import InMemoryDocumentStore
document_store = InMemoryDocumentStore()
document_store.write_documents(test_docs)
docs = document_store.get_document_ids_by_tags({'tag3': ["3"]})
assert docs == [{'name': 'testing the finder 4', 'text': 'testing the finder with pyhton unit test 3', 'meta': {'url': 'url'}, 'tags': [{'tag3': ['1', '3']}]}]