# chat-pykg/ingest.py import tempfile import gradio as gr from langchain.document_loaders import SitemapLoader, ReadTheDocsLoader, TextLoader from langchain.embeddings import OpenAIEmbeddings, HuggingFaceEmbeddings from langchain.text_splitter import RecursiveCharacterTextSplitter, PythonCodeTextSplitter, MarkdownTextSplitter, TextSplitter from langchain.vectorstores.faiss import FAISS import os from langchain.vectorstores import Chroma import shutil from pathlib import Path import subprocess import chromadb import magic from typing import Any, Dict, Iterable, List, Optional, Type, TypeVar from pydantic import Extra, Field, root_validator import logging logger = logging.getLogger() from langchain.docstore.document import Document import numpy as np def embedding_chooser(embedding_radio): if embedding_radio == "Sentence Transformers": embedding_function = HuggingFaceEmbeddings() elif embedding_radio == "OpenAI": embedding_function = OpenAIEmbeddings() else: embedding_function = HuggingFaceEmbeddings() return embedding_function # Monkeypatch pending PR def _merge_splits(self, splits: Iterable[str], separator: str) -> List[str]: # We now want to combine these smaller pieces into medium size # chunks to send to the LLM. separator_len = self._length_function(separator) docs = [] current_doc: List[str] = [] total = 0 for index, d in enumerate(splits): _len = self._length_function(d) if ( total + _len + (separator_len if len(current_doc) > 0 else 0) > self._chunk_size ): if total > self._chunk_size: logger.warning( f"Created a chunk of size {total}, " f"which is longer than the specified {self._chunk_size}" ) if len(current_doc) > 0: doc = self._join_docs(current_doc, separator) if doc is not None: docs.append(doc) # Keep on popping if: # - we have a larger chunk than in the chunk overlap # - or if we still have any chunks and the length is long while total > self._chunk_overlap or ( total + _len + (separator_len if len(current_doc) > 0 else 0) > self._chunk_size and total > 0 ): total -= self._length_function(current_doc[0]) + ( separator_len if len(current_doc) > 1 else 0 ) current_doc = current_doc[1:] if index > 0: current_doc.append(separator + d) else: current_doc.append(d) total += _len + (separator_len if len(current_doc) > 1 else 0) doc = self._join_docs(current_doc, separator) if doc is not None: docs.append(doc) return docs def get_text(content): relevant_part = content.find("div", {"class": "markdown"}) if relevant_part is not None: return relevant_part.get_text(separator=" ") else: return "" def ingest_docs(all_collections_state, urls, chunk_size, chunk_overlap, vectorstore_radio, embedding_radio, debug=False): cleared_list = urls.copy() def sanitize_folder_name(folder_name): if folder_name != '': folder_name = folder_name.strip().rstrip('/') else: folder_name = '.' # current directory return folder_name def is_hidden(path): return os.path.basename(path).startswith('.') embedding_function = embedding_chooser(embedding_radio) all_docs = [] shutil.rmtree('downloaded/', ignore_errors=True) known_exts = ["py", "md"] # Initialize text splitters py_splitter = PythonCodeTextSplitter(chunk_size=int(chunk_size), chunk_overlap=int(chunk_overlap)) text_splitter = RecursiveCharacterTextSplitter(chunk_size=int(chunk_size), chunk_overlap=int(chunk_overlap)) md_splitter = MarkdownTextSplitter(chunk_size=int(chunk_size), chunk_overlap=int(chunk_overlap)) py_splitter._merge_splits = _merge_splits.__get__(py_splitter, TextSplitter) # Process input URLs urls = [[url.strip(), [sanitize_folder_name(folder) for folder in url_folders.split(',')]] for url, url_folders in urls] for j in range(len(urls)): orgrepo = urls[j][0] repo_folders = urls[j][1] if orgrepo == '': continue if orgrepo.replace('/','-') in all_collections_state: logging.info(f"Skipping {orgrepo} as it is already in the database") continue documents_split = [] documents = [] paths = [] paths_by_ext = {} docs_by_ext = {} for ext in known_exts + ["other"]: docs_by_ext[ext] = [] paths_by_ext[ext] = [] if orgrepo[0] == '/' or orgrepo[0] == '.': # Ingest local folder local_repo_path = sanitize_folder_name(orgrepo[1:]) else: # Ingest remote git repo org = orgrepo.split('/')[0] repo = orgrepo.split('/')[1] repo_url = f"https://github.com/{org}/{repo}.git" local_repo_path = os.path.join('.downloaded', orgrepo) if debug else tempfile.mkdtemp() # Initialize the Git repository subprocess.run(["git", "init"], cwd=local_repo_path) # Add the remote repository subprocess.run(["git", "remote", "add", "-f", "origin", repo_url], cwd=local_repo_path) # Enable sparse-checkout subprocess.run(["git", "config", "core.sparseCheckout", "true"], cwd=local_repo_path) # Specify the folder to checkout cmd = ["git", "sparse-checkout", "set"] + [i for i in repo_folders] subprocess.run(cmd, cwd=local_repo_path) # Check if branch is called main or master # Checkout the desired branch res = subprocess.run(["git", "checkout", 'main'], cwd=local_repo_path) if res.returncode == 1: res = subprocess.run(["git", "checkout", "master"], cwd=local_repo_path) #res = subprocess.run(["cp", "-r", (Path(local_repo_path) / repo_folders[i]).as_posix(), '/'.join(destination.split('/')[:-1])])# # Iterate through files and process them if local_repo_path == '.': orgrepo='chat-pykg' for root, dirs, files in os.walk(local_repo_path): dirs[:] = [d for d in dirs if not is_hidden(d)] # Ignore hidden directories for file in files: if is_hidden(file): continue file_path = os.path.join(root, file) rel_file_path = os.path.relpath(file_path, local_repo_path) try: if '.' not in rel_file_path: inferred_filetype = magic.from_file(file_path, mime=True) if "python" in inferred_filetype or "text/plain" in inferred_filetype: ext = "py" else: ext = "other" else: ext = rel_file_path.split('.')[-1] if docs_by_ext.get(ext) is None: ext = "other" doc = TextLoader(os.path.join(local_repo_path, rel_file_path)).load()[0] doc.metadata["source"] = os.path.join(orgrepo, rel_file_path) docs_by_ext[ext].append(doc) paths_by_ext[ext].append(rel_file_path) except Exception as e: continue for ext in docs_by_ext.keys(): if ext == "py": documents_split += py_splitter.split_documents(docs_by_ext[ext]) documents += docs_by_ext[ext] if ext == "md": documents_split += md_splitter.split_documents(docs_by_ext[ext]) documents += docs_by_ext[ext] # else: # documents += text_splitter.split_documents(docs_by_ext[ext] all_docs += documents_split # For each document, add the metadata to the page_content for doc in documents_split: if local_repo_path != '.': doc.metadata["source"] = doc.metadata["source"].replace(local_repo_path, "") if doc.metadata["source"] == '/': doc.metadata["source"] = doc.metadata["source"][1:] doc.page_content = f'# source:{doc.metadata["source"]}\n{doc.page_content}' for doc in documents: if local_repo_path != '.': doc.metadata["source"] = doc.metadata["source"].replace(local_repo_path, "") if doc.metadata["source"] == '/': doc.metadata["source"] = doc.metadata["source"][1:] doc.page_content = f'# source:{doc.metadata["source"]}\n{doc.page_content}' if type(embedding_radio) == gr.Radio: embedding_radio = embedding_radio.value persist_directory = os.path.join(".persisted_data", embedding_radio.replace(' ','_')) persist_directory_raw = Path('.persisted_data_raw') persist_directory_raw.mkdir(parents=True, exist_ok=True) collection_name = orgrepo.replace('/','-') if vectorstore_radio == 'Chroma': collection = Chroma.from_documents(documents=documents_split, collection_name=collection_name, embedding=embedding_function, persist_directory=persist_directory) collection.persist() if vectorstore_radio == 'raw': # Persist the raw documents docarr = np.array([doc.page_content for doc in documents_split]) np.save(os.path.join(persist_directory_raw, f"{collection_name}.npy"), docarr) # with open(os.path.join(persist_directory_raw, f"{collection_name}"), "w") as f: # for doc in documents: # f.write(doc.page_content) all_collections_state.append(collection_name) cleared_list[j][0], cleared_list[j][1] = '', '' return all_collections_state, gr.update(value=cleared_list) if __name__ == "__main__": ingest_docs()