Merge pull request #3 from ittia-research/dev

add env RAG_CHUNK_SIZES

Author: etwk

.env

@@ -8,3 +8,4 @@ RAG_MODEL_DEPLOY=local
 RERANK_MODEL_NAME=BAAI/bge-reranker-v2-m3
 RERANK_BASE_URL=http://xinference:9997/v1
 SEARCH_BASE_URL=https://s.jina.ai
+RAG_CHUNK_SIZES=[4096, 1024, 256]

src/index.py

@@ -97,7 +97,7 @@ def get_contexts(statement, keywords, text):
     document = Document(text=text)
     index = build_automerging_index(
         [document],
-        chunk_sizes=[8192, 2048, 512],
+        chunk_sizes=settings.RAG_CHUNK_SIZES,
     )  # todo: will it better to use retriever directly?
 
     query_engine = get_automerging_query_engine(index, similarity_top_k=16)

src/settings.py

@@ -1,4 +1,4 @@
-import os
+import os, ast
 
 class Settings:
     def __init__(self):
@@ -15,4 +15,12 @@ def __init__(self):
         # set RAG model deploy mode
         self.RAG_MODEL_DEPLOY = os.environ.get("RAG_MODEL_DEPLOY") or "local"
 
+        # set RAG chunk sizes
+        self.RAG_CHUNK_SIZES = [4096, 1024, 256]
+        _chunk_sizes = os.environ.get("RAG_CHUNK_SIZES")
+        try:
+            self.RAG_CHUNK_SIZES = ast.literal_eval(_chunk_sizes)
+        except:
+            pass
+    
 settings = Settings()

src/utils.py

@@ -191,6 +191,8 @@ async def get_homepage():
 
 [Usage] {settings.PROJECT_HOSTING_BASE_URL}/YOUR_FACT_CHECK_QUERY
 
+[Source] https://github.com/ittia-research/check
+
 {md}
 """
     return md
@@ -201,6 +203,7 @@ async def get_stack():
         "LLM model": settings.LLM_MODEL_NAME,
         "Embedding model": settings.EMBEDDING_MODEL_NAME,
         "Rerank model": settings.RERANK_MODEL_NAME,
+        "RAG chunk sizes": settings.RAG_CHUNK_SIZES,
         "RAG deploy mode": settings.RAG_MODEL_DEPLOY,
     }
     return stack