feat: v1

constants.py

@@ -32,12 +32,12 @@ CHROMA_SETTINGS = Settings(
 )
 
 # Context Window and Max New Tokens
-CONTEXT_WINDOW_SIZE = 1024
+CONTEXT_WINDOW_SIZE = 2048
 MAX_NEW_TOKENS = CONTEXT_WINDOW_SIZE  # int(CONTEXT_WINDOW_SIZE/4)
 
 #### If you get a "not enough space in the buffer" error, you should reduce the values below, start with half of the original values and keep halving the value until the error stops appearing
 
-N_GPU_LAYERS = 50  # Llama-2-70B has 83 layers
+N_GPU_LAYERS = 83  # Llama-2-70B has 83 layers
 N_BATCH = 2048
 
 ### From experimenting with the Llama-2-7B-Chat-GGML model on 8GB VRAM, these values work:

main.py

@@ -136,28 +136,33 @@ def delete_source_route(data: Delete):
          raise HTTPException(status_code=400, detail=print(f"The file {file_to_delete} does not exist."))
 
 @api_app.post('/predict')
-async def predict(data: Predict):
+def predict(data: Predict):
     global QA
-    user_prompt = data.prompt
-    if user_prompt:
-        res = QA(user_prompt)
+    try:
+        user_prompt = data.prompt
+        if user_prompt:
+            res = QA(user_prompt)
 
-        answer, docs = res["result"], res["source_documents"]
+            answer, docs = res["result"], res["source_documents"]
 
-        prompt_response_dict = {
-            "Prompt": user_prompt,
-            "Answer": answer,
-        }
+            prompt_response_dict = {
+                "Prompt": user_prompt,
+                "Answer": answer,
+            }
+
+            prompt_response_dict["Sources"] = []
+            for document in docs:
+                prompt_response_dict["Sources"].append(
+                    (os.path.basename(str(document.metadata["source"])), str(document.page_content))
+                )
+
+            return {"response": prompt_response_dict}
+        else:
+            raise HTTPException(status_code=400, detail="Prompt Incorrect")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error occurred: {str(e)}")
 
-        prompt_response_dict["Sources"] = []
-        for document in docs:
-            prompt_response_dict["Sources"].append(
-                (os.path.basename(str(document.metadata["source"])), str(document.page_content))
-            )
 
-        return {"response": prompt_response_dict}
-    else:
-        raise HTTPException(status_code=400, detail="Prompt Incorrect")
 
 @api_app.post("/save_document/")
 async def create_upload_file(file: UploadFile):
@@ -208,12 +213,9 @@ async def websocket_endpoint(websocket: WebSocket,  client_id: int):
     try:
         while True:
             prompt = await websocket.receive_text()
-
             response = QA(inputs=prompt, return_only_outputs=True, tags=f'{client_id}', include_run_info=True)
-
             await websocket.send_text(f'{response}')
 
-
     except WebSocketDisconnect:
         print('disconnect')
     except RuntimeError as error: