Update gen_doc_chunks

app.py

@@ -85,7 +85,7 @@ from langchain.embeddings import (
 from langchain.llms import HuggingFacePipeline, OpenAI
 from langchain.memory import ConversationBufferMemory
 from langchain.text_splitter import (
-    CharacterTextSplitter,
+    # CharacterTextSplitter,
     RecursiveCharacterTextSplitter,
 )
 from langchain.vectorstores import FAISS, Chroma
@@ -97,8 +97,6 @@ from transformers import LlamaForCausalLM, LlamaTokenizer, pipeline
 from epub_loader import EpubLoader
 from load_api_key import load_api_key, pk_base, sk_base
 
-MODEL_NAME = "paraphrase-multilingual-mpnet-base-v2"  # 1.11G
-
 # fix timezone
 os.environ["TZ"] = "Asia/Shanghai"
 try:
@@ -135,6 +133,10 @@ CHROMA_SETTINGS = Settings(
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+MODEL_NAME = "paraphrase-multilingual-mpnet-base-v2"  # 1.11G
+CHUNK_SIZE = 1000  # 250
+CHUNK_OVERLAP = 100  # 50
+
 ns_initial = SimpleNamespace(
     db=None,
     qa=None,
@@ -143,8 +145,8 @@ ns_initial = SimpleNamespace(
     files_info=None,
     files_uploaded=[],
     db_ready=None,
-    chunk_size=250,
-    chunk_overlap=250,
+    chunk_size=CHUNK_SIZE,
+    chunk_overlap=CHUNK_OVERLAP,
     model_name=MODEL_NAME,
 )
 ns = deepcopy(ns_initial)
@@ -226,65 +228,94 @@ def get_pdf_text(pdf_docs):
     return text
 
 
-def get_text_chunks(text, chunk_size=1000):
-    """docs-chat."""
-    text_splitter = CharacterTextSplitter(
-        separator="\n", chunk_size=chunk_size, chunk_overlap=200, length_function=len
+# def get_text_chunks(text, chunk_size=None, chunk_overlap=None):
+def get_doc_chunks(doc: Document, chunk_size=None, chunk_overlap=None) -> List[Document]:
+    """Generate doc chunks."""
+    if chunk_size is None:
+        chunk_size = ns.chunk_size
+    if chunk_overlap is None:
+        chunk_overlap = ns.chunk_overlap
+
+    # text_splitter = CharacterTextSplitter(
+    text_splitter = RecursiveCharacterTextSplitter(
+        # separator="\n",
+        separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""],
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+        length_function=len
     )
-    chunks = text_splitter.split_text(text)
+    # chunks = text_splitter.split_text(text)
+    chunks = text_splitter.split_documents(doc)
+
     return chunks
 
 
 def get_vectorstore(
-    text_chunks,
+    # text_chunks: List[Document],
+    doc_chunks: List[Document],
     vectorstore=None,
+    model_name=None,
     persist=True,
+    persist_directory=None
 ):
     """Gne vectorstore."""
     # embedding = OpenAIEmbeddings()
     # for HuggingFaceInstructEmbeddings
-    model_name = "hkunlp/instructor-xl"
-    model_name = "hkunlp/instructor-large"
-    model_name = "hkunlp/instructor-base"
+    # model_name = "hkunlp/instructor-xl"
+    # model_name = "hkunlp/instructor-large"
+    # model_name = "hkunlp/instructor-base"
 
     # embedding = HuggingFaceInstructEmbeddings(model_name=model_name)
 
-    model_name = MODEL_NAME
+    if vectorstore is None:
+        vectorstore = "chroma"
+
+    if model_name is None:
+        model_name = MODEL_NAME
+
+    if persist_directory is None:
+        persist_directory = PERSIST_DIRECTORY
+
     logger.info(f"Loading {model_name}")
     embedding = SentenceTransformerEmbeddings(model_name=model_name)
     logger.info(f"Done loading {model_name}")
 
-    if vectorstore is None:
-        vectorstore = "chroma"
-
     if vectorstore.lower() in ["chroma"]:
         logger.info(
-            "Doing vectorstore Chroma.from_texts(texts=text_chunks, embedding=embedding)"
+            # "Doing vectorstore Chroma.from_texts(texts=text_chunks, embedding=embedding)"
+            "Doing vectorstore Chroma.from_documents(texts=doc_chunks, embedding=embedding)"
         )
         if persist:
-            vectorstore = Chroma.from_texts(
-                texts=text_chunks,
+            # vectorstore = Chroma.from_texts(
+            vectorstore = Chroma.from_documents(
+                # texts=text_chunks,
+                documents=doc_chunks,
                 embedding=embedding,
                 persist_directory=PERSIST_DIRECTORY,
                 client_settings=CHROMA_SETTINGS,
             )
         else:
-            vectorstore = Chroma.from_texts(texts=text_chunks, embedding=embedding)
+            # vectorstore = Chroma.from_texts(texts=text_chunks, embedding=embedding)
+            vectorstore = Chroma.from_documents(documents=doc_chunks, embedding=embedding)
 
         logger.info(
-            "Done vectorstore FAISS.from_texts(texts=text_chunks, embedding=embedding)"
+            # "Done vectorstore Chroma.from_texts(texts=text_chunks, embedding=embedding)"
+            "Done vectorstore Chroma.from_texts(documents=doc_chunks, embedding=embedding)"
         )
 
         return vectorstore
 
     # if vectorstore.lower() not in ['chroma']
+
     # TODO handle other cases
+
     logger.info(
-        "Doing vectorstore FAISS.from_texts(texts=text_chunks, embedding=embedding)"
+        "Doing vectorstore FAISS.from_texts(documents=doc_chunks, embedding=embedding)"
     )
-    vectorstore = FAISS.from_texts(texts=text_chunks, embedding=embedding)
+    # vectorstore = FAISS.from_texts(documents=doc_chunks, embedding=embedding)
+    vectorstore = FAISS.from_documents(documents=doc_chunks, embedding=embedding)
     logger.info(
-        "Done vectorstore FAISS.from_texts(texts=text_chunks, embedding=embedding)"
+        "Done vectorstore FAISS.from_documents(documents=doc_chunks, embedding=embedding)"
     )
 
     return vectorstore
@@ -386,11 +417,10 @@ def embed_files(progress=gr.Progress()):
     # initialize if necessary
     if ns.db is None:
         logger.info(f"loading {ns.model_name:}")
+        embedding = SentenceTransformerEmbeddings(
+            model_name=ns.model_name, model_kwargs={"device": DEVICE}
+        )
         for _ in progress.tqdm(range(1), desc="diggin..."):
-            embedding = SentenceTransformerEmbeddings(
-                model_name=ns.model_name, model_kwargs={"device": DEVICE}
-            )
-
             logger.info("creating vectorstore")
             ns.db = Chroma(
                 # persist_directory=PERSIST_DIRECTORY,

ggml-try.py

@@ -2,8 +2,9 @@
 
 https://raw.githubusercontent.com/imartinez/privateGPT/main/requirements.txt
 
--c https://gpt4all.io/models/ggml-gpt4all-j-v1.3-groovy.bin
-"""
+from pathlib import Path
+Path("models").mkdir(exit_ok=True)
+!time wget -c https://gpt4all.io/models/ggml-gpt4all-j-v1.3-groovy.bin -O models/ggml-gpt4all-j-v1.3-groovy.bin"""
 from dotenv import load_dotenv, dotenv_values
 from langchain.chains import RetrievalQA
 from langchain.embeddings import HuggingFaceEmbeddings
@@ -37,16 +38,14 @@ settings = dict([('PERSIST_DIRECTORY', 'db1'),
 
 # models/ggml-gpt4all-j-v1.3-groovy.bin  ~5G
 
-# all-MiniLM-L6-v2 () or
-embeddings_model_name = settings.get("EMBEDDINGS_MODEL_NAME")
-
-# embeddings_model_name = 'all-MiniLM-L6-v2'
-embeddings_model_name = 'paraphrase-multilingual-mpnet-base-v2'
-
 persist_directory = settings.get('PERSIST_DIRECTORY')
 
 model_type = settings.get('MODEL_TYPE')
 model_path = settings.get('MODEL_PATH')
+embeddings_model_name = settings.get("EMBEDDINGS_MODEL_NAME")
+# embeddings_model_name = 'all-MiniLM-L6-v2'
+# embeddings_model_name = 'paraphrase-multilingual-mpnet-base-v2'
+
 model_n_ctx = settings.get('MODEL_N_CTX')
 model_n_batch = int(settings.get('MODEL_N_BATCH',8))
 target_source_chunks = int(settings.get('TARGET_SOURCE_CHUNKS',4))
@@ -60,6 +59,7 @@ CHROMA_SETTINGS = Settings(
 
 args = SimpleNamespace(hide_source=False, mute_stream=False)
 
+# load chroma database from db1
 embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)
 db = Chroma(persist_directory=persist_directory, embedding_function=embeddings, client_settings=CHROMA_SETTINGS)
 
@@ -78,4 +78,19 @@ match model_type:
 
 # need about 5G RAM
 
-qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever, return_source_documents= not args.hide_source)
+qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever, return_source_documents= not args.hide_source)
+
+# Get the answer from the chain
+
+query = "共产党是什么"
+
+start = time.time()
+res = qa(query)
+answer, docs = res['result'], [] if args.hide_source else res['source_documents']
+end = time.time()
+
+# Print the result
+print("\n\n> Question:")
+print(query)
+print(f"\n> Answer (took {round(end - start, 2)} s.):")
+print(answer)