Merge pull request #57 from arjbingly/tests

Basic RAG test and DeepLake issue resolved

projects/Basic-RAG/BasicRAG_stuff.py

@@ -6,6 +6,7 @@
 
 client = DeepLakeClient(collection_name="test")
 retriever = Retriever(vectordb=client)
+
 rag = BasicRAG(doc_chain="stuff", retriever=retriever)
 
 if __name__ == "__main__":

pyproject.toml

@@ -117,3 +117,6 @@ docstring-code-format = true
 
 [tool.ruff.lint.pydocstyle]
 convention = "google"
+
+[tool.mypy]
+ignore_missing_imports = true

src/config.ini

@@ -1,5 +1,5 @@
 [llm]
-model_name : Llama-2-7b-chat
+model_name : Llama-2-13b-chat
 # meta-llama/Llama-2-70b-chat-hf Mixtral-8x7B-Instruct-v0.1
 quantization : Q5_K_M
 pipeline : llama_cpp

src/grag/components/multivec_retriever.py

@@ -49,7 +49,7 @@ def __init__(
         store_path: str = multivec_retriever_conf["store_path"],
         id_key: str = multivec_retriever_conf["id_key"],
         namespace: str = multivec_retriever_conf["namespace"],
-        top_k=1,
+        top_k=int(multivec_retriever_conf["top_k"]),
         client_kwargs: Optional[Dict[str, Any]] = None,
     ):
         """Initialize the Retriever.

src/tests/rag/basic_rag_test.py

@@ -1,21 +1,25 @@
-from typing import Text, List
+from typing import List, Text
 
+from grag.components.multivec_retriever import Retriever
+from grag.components.vectordb.deeplake_client import DeepLakeClient
 from grag.rag.basic_rag import BasicRAG
 
+client = DeepLakeClient(collection_name="test")
+retriever = Retriever(vectordb=client)
+
 
 def test_rag_stuff():
-    rag = BasicRAG(doc_chain="stuff")
-    response, sources = rag("What is simulated annealing?")
+    rag = BasicRAG(doc_chain="stuff", retriever=retriever)
+    response, sources = rag("What is Flash Attention?")
     assert isinstance(response, Text)
     assert isinstance(sources, List)
     assert all(isinstance(s, str) for s in sources)
     del rag.llm
 
 
 def test_rag_refine():
-    rag = BasicRAG(doc_chain="refine")
-    response, sources = rag("What is simulated annealing?")
-    # assert isinstance(response, Text)
+    rag = BasicRAG(doc_chain="refine", retriever=retriever)
+    response, sources = rag("What is Flash Attention?")
     assert isinstance(response, List)
     assert all(isinstance(s, str) for s in response)
     assert isinstance(sources, List)