Added cli commands

scripts/ingest.py

@@ -5,12 +5,8 @@ import dotenv
 import typer
 
 from collections import defaultdict
-from pathlib import Path
 from typing import List, Optional
 
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-
-
 from parser.file.bulk import SimpleDirectoryReader
 from parser.schema.base import Document
 from parser.open_ai_func import call_openai_api, get_user_permission
@@ -39,14 +35,17 @@ def ingest(yes: bool = typer.Option(False, "-y", "--yes", prompt=False,
            file: Optional[List[str]] = typer.Option(None,
                                                    help="""File paths to use (Optional; overrides dir).
                                                         E.g. --file inputs/1.md --file inputs/2.md"""),
-           recursive: Optional[bool] = typer.Option(True,
-                                                   help="Whether to recursively search in subdirectories."),
-           limit: Optional[int] = typer.Option(None,
-                                                   help="Maximum number of files to read."),
+           recursive: Optional[bool] = typer.Option(True, help="Whether to recursively search in subdirectories."),
+           limit: Optional[int] = typer.Option(None, help="Maximum number of files to read."),
            formats: Optional[List[str]] = typer.Option([".rst", ".md"],
                                                    help="""List of required extensions (list with .)
                                                         Currently supported: .rst, .md, .pdf, .docx, .csv, .epub, .html, .mdx"""),
-           exclude: Optional[bool] = typer.Option(True, help="Whether to exclude hidden files (dotfiles).")):
+           exclude: Optional[bool] = typer.Option(True, help="Whether to exclude hidden files (dotfiles)."),
+           sample: Optional[bool] = typer.Option(False, help="Whether to output sample of the first 5 split documents."),
+           token_check: Optional[bool] = typer.Option(True, help="Whether to group small documents and split large."),
+           min_tokens: Optional[int] = typer.Option(150, help="Minimum number of tokens to not group."),
+           max_tokens: Optional[int] = typer.Option(2000, help="Maximum number of tokens to not split."),
+           ):
 
     """
         Creates index from specified location or files.
@@ -57,16 +56,22 @@ def ingest(yes: bool = typer.Option(False, "-y", "--yes", prompt=False,
         raw_docs = SimpleDirectoryReader(input_dir=directory, input_files=file, recursive=recursive,
                                          required_exts=formats, num_files_limit=limit,
                                          exclude_hidden=exclude).load_data()
-        #Checking min_tokens and max_tokens
-        raw_docs = group_split(documents=raw_docs)
 
-        docs = [Document.to_langchain_format(raw_doc) for raw_doc in raw_docs]
         # Here we split the documents, as needed, into smaller chunks.
         # We do this due to the context limits of the LLMs.
-
+        raw_docs = group_split(documents=raw_docs, min_tokens=min_tokens, max_tokens=max_tokens, token_check=token_check)
+        #Old method
         # text_splitter = RecursiveCharacterTextSplitter()
         # docs = text_splitter.split_documents(raw_docs)
 
+        #Sample feature
+        if sample == True:
+            for i in range(min(5, len(raw_docs))):
+                print(raw_docs[i].text)
+
+        docs = [Document.to_langchain_format(raw_doc) for raw_doc in raw_docs]
+
+
         # Here we check for command line arguments for bot calls.
         # If no argument exists or the yes is not True, then the
         # user permission is requested to call the API.