File size: 4,824 Bytes
60d3c71 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | import copy
import typing
from nltk.corpus import stopwords
import nltk
nltk.download('stopwords')
STOP= stopwords.words('english')
from nltk.tokenize import word_tokenize
import string
from cherche import retrieve
from lenlp import sparse
__all__ = ["Retriever"]
class Retriever:
"""Knowledge retriever.
Parameters
----------
documents
List of documents.
Examples:
---------
>>> import json
>>> from knowledge_database import retriever
>>> with open("database/database.json", "r") as f:
... documents = json.load(f)
>>> knowledge_retriever = retriever.Retriever(documents=documents)
>>> candidates = knowledge_retriever.documents("neural search")
>>> candidates = knowledge_retriever.tags("neural search")
>>> candidates = knowledge_retriever.documents_tags("neural search")
"""
def __init__(self, documents: typing.Dict):
updated_documents = copy.deepcopy(documents)
for key,_ in updated_documents.items():
summary= updated_documents[key]["summary"]
title = updated_documents[key]["title"]
summary_ready=[]
title_ready=[]
summary_words=word_tokenize(summary.lower())
title_words=word_tokenize(title.lower())
for i in summary_words:
if (i not in STOP ) and (i not in string.punctuation ) and (not i.isdigit()) :
summary_ready.append(i)
for i in title_words:
if (i not in STOP ) and (i not in string.punctuation ) and (not i.isdigit()) :
title_ready.append(i)
updated_documents[key]["summary"]=" ".join(summary_ready)
updated_documents[key]['title']=" ".join(title_ready)
documents = [{"url": url, **document} for url, document in documents.items()]
updated_documents = [
{
**{
"url": url,
"tags": " ".join(document.pop("tags") + document.pop("extra-tags")),
},
**document,
}
for url, document in updated_documents.items()
]
self.retriever = (
retrieve.TfIdf(
key="url",
on=["title", "tags", "summary", "date"],
k=30,
tfidf=sparse.BM25Vectorizer(
normalize=True,
ngram_range=(4, 7),
analyzer="char_wb",
b=0,
),
documents=updated_documents,
)
| retrieve.TfIdf(
key="url",
on=["title", "tags", "summary", "date"],
k=10,
tfidf=sparse.BM25Vectorizer(
normalize=True,
ngram_range=(2, 5),
analyzer="char_wb",
),
documents=updated_documents,
)
) + documents
# Retrieve documents that match a specific tag.
self.retriever_documents_tags = (
retrieve.TfIdf(
key="url",
on=["title", "tags", "summary", "date"],
k=40,
tfidf=sparse.BM25Vectorizer(
normalize=True,
ngram_range=(4, 7),
analyzer="char_wb",
),
documents=updated_documents,
)
& retrieve.TfIdf(
key="url",
on=["tags"],
k=40,
tfidf=sparse.BM25Vectorizer(
normalize=True,
ngram_range=(4, 7),
analyzer="char_wb",
),
documents=updated_documents,
)
) + documents
tags = {}
for document in documents:
for tag in document["tags"] + document["extra-tags"]:
tags[tag] = True
tags = [{"tag": tag} for tag in tags]
self.retriever_tags = (
retrieve.TfIdf(
key="tag",
on=["tag"],
k=5,
tfidf=sparse.BM25Vectorizer(
normalize=True,
ngram_range=(3, 7),
analyzer="char_wb",
),
documents=tags,
)
+ tags
)
def documents(self, q: str):
"""Match documents."""
return self.retriever(q)
def tags(self, q: str):
"""Match tags."""
return [tag["tag"] for tag in self.retriever_tags(q)]
def documents_tags(self, q: str):
"""Match documents and tags."""
return self.retriever_documents_tags(q)
|