(self)
| 26 | |
| 27 | class WikipediaIndex: |
| 28 | def __init__(self): |
| 29 | path = os.path.join(DIR, '..', 'data', 'wikipedia-3sentence-level-retrieval-index', 'files') |
| 30 | indexpath = os.path.join(path, 'knn.index') |
| 31 | wiki_sentence_path = os.path.join(path, 'wikipedia-en-sentences.parquet') |
| 32 | |
| 33 | self._device = 'cuda' |
| 34 | self._tokenizer = AutoTokenizer.from_pretrained('facebook/contriever-msmarco') |
| 35 | self._contriever = AutoModel.from_pretrained('facebook/contriever-msmarco').to(self._device) |
| 36 | |
| 37 | self._df_sentences = pd.read_parquet(wiki_sentence_path, engine='fastparquet') |
| 38 | |
| 39 | self._wiki_index = faiss.read_index(indexpath, faiss.IO_FLAG_MMAP | faiss.IO_FLAG_READ_ONLY) |
| 40 | |
| 41 | |
| 42 | def search(self, query, k=1, w=5, w_th=0.5): |
nothing calls this directly
no test coverage detected