feat: add websocket

ingest.py

@@ -124,7 +124,7 @@ def main(device_type):
     text_documents, python_documents = split_documents(documents)
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
     python_splitter = RecursiveCharacterTextSplitter.from_language(
-        language=Language.PYTHON, chunk_size=880, chunk_overlap=200
+        language=Language.PYTHON, chunk_size=1000, chunk_overlap=200
     )
     texts = text_splitter.split_documents(text_documents)
     texts.extend(python_splitter.split_documents(python_documents))
@@ -151,6 +151,9 @@ def main(device_type):
 
     )
 
+    db.persist()
+    db = None
+
 if __name__ == "__main__":
     logging.basicConfig(
         format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)s - %(message)s", level=logging.INFO

load_models.py

@@ -58,13 +58,13 @@ def load_quantized_model_gguf_ggml(model_id, model_basename, device_type, loggin
             "model_path": model_path,
             "n_ctx": CONTEXT_WINDOW_SIZE,
             "max_tokens": MAX_NEW_TOKENS,
-            "n_batch": N_BATCH,
              # set this based on your GPU & CPU RAM
         }
         if device_type.lower() == "mps":
             kwargs["n_gpu_layers"] = 1
         if device_type.lower() == "cuda":
-            kwargs["n_gpu_layers"] = N_GPU_LAYERS  # set this based on your GPU
+            kwargs["n_gpu_layers"] = N_GPU_LAYERS
+            kwargs["n_batch"] = MAX_NEW_TOKENS  # set this based on your GPU
 
         kwargs["stream"] = stream
 

main.py

@@ -56,8 +56,6 @@ QA = RetrievalQA.from_chain_type(
 )
 
 def sendPromptChain(QA, user_prompt):
-    print(QA)
-
     res = QA(user_prompt)
 
     answer, docs = res["result"], res["source_documents"]
@@ -290,7 +288,7 @@ async def websocket_endpoint_room(websocket: WebSocket, room_id: str, user_id: s
         while True:
             data = await websocket.receive_text()
 
-            prompt_response_dict = sendPromptChain(QA, data, True)
+            prompt_response_dict = sendPromptChain(QA, data)
 
             await socket_manager.broadcast_to_room(room_id, json.dumps(prompt_response_dict))