performance optimization and bug fix

app.py

@@ -11,9 +11,6 @@ import transformers
 from torch import cuda, bfloat16 
 
 
-
-load_dotenv()
-
 embeddings_model_name = "all-MiniLM-L6-v2"
 persist_directory = "db"
 model = "tiiuae/falcon-7b-instruct"
@@ -26,6 +23,10 @@ source_directory = os.environ.get('SOURCE_DIRECTORY', 'source_documents')
 
 from constants import CHROMA_SETTINGS
 
+
+embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)
+llm = HuggingFacePipeline.from_model_id(model_id=model, task="text-generation", device=0, model_kwargs={"temperature":0.1,"trust_remote_code": True, "max_length":100000, "top_p":0.15, "top_k":0, "repetition_penalty":1.1, "num_return_sequences":1,})
+
 # async def test_embedding():
 #     # Create the folder if it doesn't exist
 #     os.makedirs(source_directory, exist_ok=True)
@@ -101,14 +102,10 @@ def embed_documents(files, collection_name: Optional[str] = None):
 def retrieve_documents(query: str, collection_name:str):
     target_source_chunks = 4
     mute_stream = ""
-    embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)
     db = Chroma(persist_directory=persist_directory,collection_name=collection_name, embedding_function=embeddings, client_settings=CHROMA_SETTINGS)
     retriever = db.as_retriever(search_kwargs={"k": target_source_chunks})
     # Prepare the LLM
     callbacks = [] if mute_stream else [StreamingStdOutCallbackHandler()]
-
-    llm = HuggingFacePipeline.from_model_id(model_id=model, task="text-generation", device=0, model_kwargs={"temperature":0.1,"trust_remote_code": True, "max_length":100000, "top_p":0.15, "top_k":0, "repetition_penalty":1.1, "num_return_sequences":1,})
-
     qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever, return_source_documents=False)
     
     # Get the answer from the chain

privateGPT.py

@@ -7,12 +7,12 @@ import os
 
 load_dotenv()
 
-embeddings_model_name = os.environ.get("EMBEDDINGS_MODEL_NAME")
-persist_directory = os.environ.get('PERSIST_DIRECTORY')
+# embeddings_model_name = os.environ.get("EMBEDDINGS_MODEL_NAME")
+# persist_directory = os.environ.get('PERSIST_DIRECTORY')
 
-model_type = os.environ.get('MODEL_TYPE')
-model_path = os.environ.get('MODEL_PATH')
-model_n_ctx = os.environ.get('MODEL_N_CTX')
+# model_type = os.environ.get('MODEL_TYPE')
+# model_path = os.environ.get('MODEL_PATH')
+# model_n_ctx = os.environ.get('MODEL_N_CTX')
 
 from constants import CHROMA_SETTINGS
 

streamlit_app.py

@@ -7,7 +7,7 @@ import socket
 from urllib3.connection import HTTPConnection
 from app import embed_documents, retrieve_documents
 
-API_BASE_URL = os.environ.get("API_BASE_URL")
+# API_BASE_URL = os.environ.get("API_BASE_URL")
 
 
 embeddings_model_name = "all-MiniLM-L6-v2"
@@ -86,7 +86,3 @@ def get_collection_names():
 #     else:
 #         st.error("Failed to retrieve documents.")
 #         st.write(response.text)
-
-
-if __name__ == "__main__":
-    main()