2020-04-28 07:10:32 -07:00
|
|
|
def test_tfidf_retriever():
|
2020-06-30 19:05:45 +02:00
|
|
|
from haystack.retriever.sparse import TfidfRetriever
|
2020-04-28 07:10:32 -07:00
|
|
|
|
|
|
|
test_docs = [
|
2020-07-31 11:34:06 +02:00
|
|
|
{"id": "26f84672c6d7aaeb8e2cd53e9c62d62d", "name": "testing the finder 1", "text": "godzilla says hello"},
|
2020-04-28 07:10:32 -07:00
|
|
|
{"name": "testing the finder 2", "text": "optimus prime says bye"},
|
|
|
|
{"name": "testing the finder 3", "text": "alien says arghh"}
|
|
|
|
]
|
|
|
|
|
|
|
|
from haystack.database.memory import InMemoryDocumentStore
|
|
|
|
document_store = InMemoryDocumentStore()
|
|
|
|
document_store.write_documents(test_docs)
|
|
|
|
|
|
|
|
retriever = TfidfRetriever(document_store)
|
|
|
|
retriever.fit()
|
2020-07-31 11:34:06 +02:00
|
|
|
doc = retriever.retrieve("godzilla", top_k=1)[0]
|
2020-08-03 16:20:17 +02:00
|
|
|
assert doc.id == "26f84672c6d7aaeb8e2cd53e9c62d62d"
|
2020-07-31 11:34:06 +02:00
|
|
|
assert doc.text == 'godzilla says hello'
|
|
|
|
assert doc.meta == {"name": "testing the finder 1"}
|