| import shutil |
| import time |
|
|
| import lancedb |
| import openai |
| import pyarrow as pa |
| import pandas as pd |
| from pathlib import Path |
| import tqdm |
| import numpy as np |
|
|
| from gradio_app.backend.embedders import EmbedderFactory |
| from markdown_to_text import * |
| from settings import * |
|
|
|
|
| with open('data/openaikey.txt') as f: |
| OPENAI_KEY = f.read().strip() |
| openai.api_key = OPENAI_KEY |
|
|
|
|
| |
| db = lancedb.connect(LANCEDB_DIRECTORY) |
| batch_size = 32 |
|
|
| schema = pa.schema([ |
| pa.field(VECTOR_COLUMN_NAME, pa.list_(pa.float32(), emb_sizes[EMBED_NAME])), |
| pa.field(TEXT_COLUMN_NAME, pa.string()), |
| pa.field(DOCUMENT_PATH_COLUMN_NAME, pa.string()), |
| ]) |
| table_name = f'{LANCEDB_TABLE_NAME}_{CHUNK_POLICY}_{EMBED_NAME}' |
| tbl = db.create_table(table_name, schema=schema, mode="overwrite") |
|
|
| input_dir = Path(MARKDOWN_SOURCE_DIR) |
| files = list(input_dir.rglob("*")) |
|
|
| chunks = [] |
| for file in files: |
| if not os.path.isfile(file): |
| continue |
|
|
| file_path, file_ext = os.path.splitext(os.path.relpath(file, input_dir)) |
| if file_ext != '.md': |
| print(f'Skipped {file_ext} extension: {file}') |
| continue |
|
|
| with open(file, encoding='utf-8') as f: |
| f = f.read() |
| f = remove_comments(f) |
| if CHUNK_POLICY == "txt": |
| f = md2txt_then_split(f) |
| else: |
| assert CHUNK_POLICY == "md" |
| f = split_markdown(f) |
| chunks.extend((chunk, os.path.abspath(file)) for chunk in f) |
|
|
| from matplotlib import pyplot as plt |
| plt.hist([len(c) for c, d in chunks], bins=100) |
| plt.title(table_name) |
| plt.show() |
|
|
| embedder = EmbedderFactory.get_embedder(EMBED_NAME) |
|
|
| time_embed, time_ingest = [], [] |
| for i in tqdm.tqdm(range(0, int(np.ceil(len(chunks) / batch_size)))): |
| texts, doc_paths = [], [] |
| for text, doc_path in chunks[i * batch_size:(i + 1) * batch_size]: |
| if len(text) > 0: |
| texts.append(text) |
| doc_paths.append(doc_path) |
|
|
| t = time.perf_counter() |
| encoded = embedder.embed(texts) |
| time_embed.append(time.perf_counter() - t) |
|
|
| df = pd.DataFrame({ |
| VECTOR_COLUMN_NAME: encoded, |
| TEXT_COLUMN_NAME: texts, |
| DOCUMENT_PATH_COLUMN_NAME: doc_paths, |
| }) |
|
|
| t = time.perf_counter() |
| tbl.add(df) |
| time_ingest.append(time.perf_counter() - t) |
|
|
|
|
| time_embed = sum(time_embed) |
| time_ingest = sum(time_ingest) |
| print(f'Embedding: {time_embed}, Ingesting: {time_ingest}') |
|
|
|
|
|
|
|
|
