refactor

chunking_utils.py

@@ -0,0 +1,42 @@
+import json
+import numpy as np
+
+from tqdm import tqdm
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+from models import env_config
+
+
+class Chunker:
+    def __init__(self, strategy, split_seq=".", chunk_len=512):
+        self.split_seq = split_seq
+        self.chunk_len = chunk_len
+        if strategy == "recursive":
+            self.split = RecursiveCharacterTextSplitter(
+                chunk_size=chunk_len,
+                separators=[split_seq]
+            ).split_text
+        if strategy == "sequence":
+            self.split = self.seq_splitter
+        if strategy == "constant":
+            self.split = self.const_splitter
+
+    def seq_splitter(self, text):
+        return text.split(self.split_seq)
+
+    def const_splitter(self, text):
+        return [
+            text[i * self.chunk_len:(i + 1) * self.chunk_len]
+            for i in range(int(np.ceil(len(text) / self.chunk_len)))
+        ]
+
+
+def chunk_generator(input_dataset, chunker, tmp_file):
+    for i in tqdm(range(len(input_dataset))):
+        chunks = chunker.split(input_dataset[i][env_config.input_text_col])
+        for chunk in chunks:
+            if chunk:
+                tmp_file.write(
+                    json.dumps({env_config.input_text_col: chunk}) + "\n"
+                )
+                yield {env_config.input_text_col: chunk}

embed_utils.py

@@ -0,0 +1,55 @@
+import json
+import asyncio
+import logging
+import time
+
+from tqdm.asyncio import tqdm_asyncio
+from huggingface_hub import get_inference_endpoint
+
+from models import env_config, embed_config
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+endpoint = get_inference_endpoint(env_config.tei_name, token=env_config.hf_token)
+
+
+async def embed_chunk(sentence, semaphore, tmp_file):
+    async with semaphore:
+        payload = {
+            "inputs": sentence,
+            "truncate": True
+        }
+
+        try:
+            resp = await endpoint.async_client.post(json=payload)
+        except Exception as e:
+            raise RuntimeError(str(e))
+
+        result = json.loads(resp)
+        tmp_file.write(
+            json.dumps({"vector": result[0], env_config.input_text_col: sentence}) + "\n"
+        )
+
+
+async def embed_wrapper(input_ds, temp_file):
+    semaphore = asyncio.BoundedSemaphore(embed_config.semaphore_bound)
+    jobs = [
+        asyncio.create_task(embed_chunk(row[env_config.input_text_col], semaphore, temp_file))
+        for row in input_ds if row[env_config.input_text_col].strip()
+    ]
+    logger.info(f"num chunks to embed: {len(jobs)}")
+
+    tic = time.time()
+    await tqdm_asyncio.gather(*jobs)
+    logger.info(f"embed time: {time.time() - tic}")
+
+
+def wake_up_endpoint():
+    endpoint.fetch()
+    if endpoint.status != 'running':
+        logger.info("Starting up TEI endpoint")
+        endpoint.resume()
+        endpoint.wait()
+    logger.info("TEI endpoint is up")
+    return

main.py

@@ -1,46 +1,32 @@
 import asyncio
 import logging
-import numpy as np
-import time
-import json
-import os
 import tempfile
-import requests
 
 from fastapi import FastAPI, Request, BackgroundTasks
 from fastapi.responses import HTMLResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
-
-from aiohttp import ClientSession
-from langchain.text_splitter import RecursiveCharacterTextSplitter
+from contextlib import asynccontextmanager
 from datasets import Dataset, load_dataset
-from tqdm import tqdm
-from tqdm.asyncio import tqdm_asyncio
 
-from models import chunk_config, embed_config, WebhookPayload
+from models import chunk_config, embed_config, env_config, WebhookPayload
+from chunking_utils import Chunker, chunk_generator
+from embed_utils import wake_up_endpoint, embed_wrapper
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# you token from Settings
-HF_TOKEN = os.getenv("HF_TOKEN")
+app_state = {}
+
 
-# URL of TEI endpoint
-TEI_URL = os.getenv("TEI_URL")
-# name of chunked dataset
-CHUNKED_DS_NAME = os.getenv("CHUNKED_DS_NAME")
-# name of embeddings dataset
-EMBED_DS_NAME = os.getenv("EMBED_DS_NAME")
-# splits of input dataset to process, comma separated
-INPUT_SPLITS = os.getenv("INPUT_SPLITS")
-# name of column to load from input dataset
-INPUT_TEXT_COL = os.getenv("INPUT_TEXT_COL")
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    app_state["seen_Sha"] = set()
+    yield
+    app_state.clear()
 
-INPUT_SPLITS = [spl.strip() for spl in INPUT_SPLITS.split(",") if spl]
 
-app = FastAPI()
-app.state.seen_Sha = set()
+app = FastAPI(lifespan=lifespan)
 
 app.mount("/static", StaticFiles(directory="static"), name="static")
 templates = Jinja2Templates(directory="templates")
@@ -61,151 +47,68 @@ async def post_webhook(
         and payload.event.scope.startswith("repo.content")
         and payload.repo.type == "dataset"
         # webhook posts multiple requests with the same update, this addresses that
-        and payload.repo.headSha not in app.state.seen_Sha
+        and payload.repo.headSha not in app_state["seen_Sha"]
     ):
-        # no-op
         logger.info("Update detected, no action taken")
         return {"processed": False}
 
-    app.state.seen_Sha.add(payload.repo.headSha)
-    task_queue.add_task(chunk_dataset, ds_name=payload.repo.name)
-    task_queue.add_task(embed_dataset, ds_name=CHUNKED_DS_NAME)
+    app_state["seen_Sha"].add(payload.repo.headSha)
+    task_queue.add_task(chunk_and_embed, input_ds_name=payload.repo.name)
 
     return {"processed": True}
 
 
-"""
-CHUNKING
-"""
-
-class Chunker:
-    def __init__(self, strategy, split_seq=".", chunk_len=512):
-        self.split_seq = split_seq
-        self.chunk_len = chunk_len
-        if strategy == "recursive":
-            self.split = RecursiveCharacterTextSplitter(
-                chunk_size=chunk_len,
-                separators=[split_seq]
-            ).split_text
-        if strategy == "sequence":
-            self.split = self.seq_splitter
-        if strategy == "constant":
-            self.split = self.const_splitter
-
-    def seq_splitter(self, text):
-        return text.split(self.split_seq)
-
-    def const_splitter(self, text):
-        return [
-            text[i * self.chunk_len:(i + 1) * self.chunk_len]
-            for i in range(int(np.ceil(len(text) / self.chunk_len)))
-        ]
-
-
-def chunk_generator(input_dataset, chunker):
-    for i in tqdm(range(len(input_dataset))):
-        chunks = chunker.split(input_dataset[i][INPUT_TEXT_COL])
-        for chunk in chunks:
-            if chunk:
-                yield {INPUT_TEXT_COL: chunk}
-
-
-def chunk_dataset(ds_name):
+def chunk(ds_name):
     logger.info("Update detected, chunking is scheduled")
-    input_ds = load_dataset(ds_name, split="+".join(INPUT_SPLITS))
+    input_ds = load_dataset(ds_name, split="+".join(env_config.input_splits))
     chunker = Chunker(
         strategy=chunk_config.strategy,
         split_seq=chunk_config.split_seq,
         chunk_len=chunk_config.chunk_len
     )
-
+    tmp_file = tempfile.NamedTemporaryFile(mode="a", suffix=".jsonl")
     dataset = Dataset.from_generator(
         chunk_generator,
         gen_kwargs={
             "input_dataset": input_ds,
-            "chunker": chunker
+            "chunker": chunker,
+            "tmp_file": tmp_file
         }
     )
 
     dataset.push_to_hub(
-        CHUNKED_DS_NAME,
+        env_config.chunked_ds_name,
         private=chunk_config.private,
-        token=HF_TOKEN
+        token=env_config.hf_token
     )
 
     logger.info("Done chunking")
-    return {"processed": True}
-
-
-"""
-EMBEDDING
-"""
+    return tmp_file
 
-async def embed_sent(sentence, semaphore, tmp_file):
-    async with semaphore:
-        payload = {
-            "inputs": sentence,
-            "truncate": True
-        }
 
-        async with ClientSession(
-            headers={
-                "Content-Type": "application/json",
-                "Authorization": f"Bearer {HF_TOKEN}"
-            }
-        ) as session:
-            async with session.post(TEI_URL, json=payload) as resp:
-                if resp.status != 200:
-                    raise RuntimeError(await resp.text())
-                result = await resp.json()
-
-                tmp_file.write(
-                    json.dumps({"vector": result[0], INPUT_TEXT_COL: sentence}) + "\n"
-                )
-
-
-async def embed(input_ds, temp_file):
-    semaphore = asyncio.BoundedSemaphore(embed_config.semaphore_bound)
-    jobs = [
-        asyncio.create_task(embed_sent(row[INPUT_TEXT_COL], semaphore, temp_file))
-        for row in input_ds if row[INPUT_TEXT_COL].strip()
-    ]
-    logger.info(f"num chunks to embed: {len(jobs)}")
-
-    tic = time.time()
-    await tqdm_asyncio.gather(*jobs)
-    logger.info(f"embed time: {time.time() - tic}")
-
-
-def wake_up_endpoint(url):
-    logger.info("Starting up TEI endpoint")
-    n_loop = 0
-    while requests.get(
-        url=url,
-        headers={"Authorization": f"Bearer {HF_TOKEN}"}
-    ).status_code != 200:
-        time.sleep(2)
-        n_loop += 1
-        if n_loop > 40:
-            raise TimeoutError("TEI endpoint is unavailable")
-    logger.info("TEI endpoint is up")
-
-
-def embed_dataset(ds_name):
+def embed(chunked_file):
     logger.info("Update detected, embedding is scheduled")
-    wake_up_endpoint(TEI_URL)
-    input_ds = load_dataset(ds_name, split="train")
+    wake_up_endpoint()
+    chunked_ds = Dataset.from_json(chunked_file.name)
     with tempfile.NamedTemporaryFile(mode="a", suffix=".jsonl") as temp_file:
-        asyncio.run(embed(input_ds, temp_file))
+        asyncio.run(embed_wrapper(chunked_ds, temp_file))
 
-        dataset = Dataset.from_json(temp_file.name)
-        dataset.push_to_hub(
-            EMBED_DS_NAME,
+        emb_ds = Dataset.from_json(temp_file.name)
+        emb_ds.push_to_hub(
+            env_config.embed_ds_name,
             private=embed_config.private,
-            token=HF_TOKEN
+            token=env_config.hf_token
         )
 
+    chunked_file.close()
     logger.info("Done embedding")
+    return
+
+
+def chunk_and_embed(input_ds_name):
+    chunked_tmp_file = chunk(input_ds_name)
+    embed(chunked_tmp_file)
+
     return {"processed": True}
 
 

models.py

@@ -4,6 +4,21 @@ from pydantic import BaseModel
 from typing import Literal
 
 
+class EnvConfig(BaseModel):
+	# you token from Settings
+	hf_token: str = os.getenv("HF_TOKEN")
+	# NAME of TEI endpoint
+	tei_name: str = os.getenv("TEI_NAME")
+	# name of chunked dataset
+	chunked_ds_name: str = os.getenv("CHUNKED_DS_NAME")
+	# name of embeddings dataset
+	embed_ds_name: str = os.getenv("EMBED_DS_NAME")
+	# splits of input dataset to process, comma separated
+	input_splits: str = os.getenv("INPUT_SPLITS")
+	# name of column to load from input dataset
+	input_text_col: str = os.getenv("INPUT_TEXT_COL")
+
+
 class ChunkConfig(BaseModel):
 	strategy: Literal["recursive", "sequence", "constant"]
 	split_seq: str
@@ -41,3 +56,7 @@ with open(os.path.join(os.getcwd(), "configs/chunk_config.json")) as c:
 with open(os.path.join(os.getcwd(), "configs/embed_config.json")) as c:
 	data = json.load(c)
 	embed_config = EmbedConfig.model_validate_json(json.dumps(data))
+
+
+env_config = EnvConfig()
+env_config.input_splits = [spl.strip() for spl in env_config.input_splits.split(",") if spl]