simple text data sources (#19)

Signed-off-by: Max Pumperla <[email protected]>

.gitignore

@@ -107,3 +107,6 @@ airflow/airflow.db
 
 # scraped folders
 docs.ray.io/
+
+# book and other source folders
+data/

app/index.py

@@ -87,7 +87,7 @@ def path_to_uri(path, scheme="https://", domain="docs.ray.io"):
  return scheme + domain + path.split(domain)[-1]
 
 
-def parse_file(record):
+def parse_html_file(record):
  html_content = load_html_file(record["path"])
  if not html_content:
  return []
@@ -100,6 +100,17 @@ def parse_file(record):
  ]
 
 
+def parse_text_file(record):
+ with open(record["path"]) as f:
+ text = f.read()
+ return [
+ {
+ "source": str(record["path"]),
+ "text": text,
+ }
+ ]
+
+
 class EmbedChunks:
  def __init__(self, model_name):
  self.embedding_model = HuggingFaceEmbeddings(
@@ -139,6 +150,7 @@ def __call__(self, batch):
 @app.command()
 def create_index(
  docs_path: Annotated[str, typer.Option(help="location of data")] = DOCS_PATH,
+ extension_type: Annotated[str, typer.Option(help="type of data")] = "html",
  embedding_model: Annotated[str, typer.Option(help="embedder")] = EMBEDDING_MODEL,
  chunk_size: Annotated[int, typer.Option(help="chunk size")] = CHUNK_SIZE,
  chunk_overlap: Annotated[int, typer.Option(help="chunk overlap")] = CHUNK_OVERLAP,
@@ -148,11 +160,17 @@ def create_index(
 
  # Dataset
  ds = ray.data.from_items(
- [{"path": path} for path in Path(docs_path).rglob("*.html") if not path.is_dir()]
+ [
+ {"path": path}
+ for path in Path(docs_path).rglob(f"*.{extension_type}")
+ if not path.is_dir()
+ ]
  )
 
  # Sections
- sections_ds = ds.flat_map(parse_file)
+ parser = parse_html_file if extension_type == "html" else parse_text_file
+ sections_ds = ds.flat_map(parser)
+ # TODO: do we really need to take_all()? Bring the splitter to the cluster
  sections = sections_ds.take_all()
 
  # Chunking

dashboard/pages/1_✨_Generation.py

@@ -9,7 +9,7 @@
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from pgvector.psycopg import register_vector
 
-from app.index import parse_file
+from app.index import parse_html_file
 from app.query import generate_response
 
 
@@ -38,7 +38,7 @@ def get_ds(docs_path):
 docs_page_url = st.text_input("Docs page URL", "https://docs.ray.io/en/master/train/faq.html")
 docs_page_path = docs_path_str + docs_page_url.split("docs.ray.io/en/master/")[-1]
 with st.expander("View sections"):
- sections = parse_file({"path": docs_page_path})
+ sections = parse_html_file({"path": docs_page_path})
  st.write(sections)
 
 # Chunks