From db26fe20afff2e7f09767ecead512de63e739057 Mon Sep 17 00:00:00 2001 From: geemi725 Date: Wed, 28 Feb 2024 12:06:58 +0100 Subject: [PATCH] hf embedding --- tests/test_xpertai.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_xpertai.py b/tests/test_xpertai.py index b37eb0f..1e33a67 100644 --- a/tests/test_xpertai.py +++ b/tests/test_xpertai.py @@ -35,9 +35,12 @@ def test_load_split_docs(): assert isinstance(docs_split[0].page_content, str) def test_vector_db(): + from langchain.embeddings import HuggingFaceEmbeddings + base_dir = os.path.dirname(os.path.realpath("__file__")) lit_file = os.path.join(base_dir, "tests/paper_test.pdf") - vector_db(lit_file=lit_file, try_meta_data=True, clean=True) + embedding = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2") + vector_db(lit_file=lit_file, try_meta_data=True, clean=True, embedding=embedding) assert os.path.exists("./data/chroma/") is True