feat: chunk retrieval updated

src/chatbot/embedding.py

@@ -1,100 +1,71 @@
-# to run this file you need model.onnx_data on the assets/onnx folder or you can obtain it from here.: https://huggingface.co/onnx-community/embeddinggemma-300m-ONNX/tree/main/onnx
-# model can also be loaded directly from autoModel.pretrained by using the same link "onnx-community/embeddinggemma-300m-ONNX"
-
-import asyncio
 import os
+import numpy as np
 from typing import List
+import onnxruntime as ort
+from transformers import AutoTokenizer
+from huggingface_hub import hf_hub_download
 
-import numpy as np
+MODEL_ID = "onnx-community/embeddinggemma-300m-ONNX"
 
-# import onnxruntime as ort
-from transformers import AutoTokenizer
+class EmbeddingModel:
+    def __init__(self):
+        print("🔵 Loading tokenizer…")
+        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
-BASE_DIR = os.path.dirname(__file__)
+        print("🔵 Downloading ONNX model files…")
+
+        self.model_path = hf_hub_download(
+            repo_id=MODEL_ID,
+            filename="onnx/model.onnx"
+        )
+        self.data_path = hf_hub_download(
+            repo_id=MODEL_ID,
+            filename="onnx/model.onnx_data"
+        )
 
-# TOKENIZER_DIR = os.path.abspath(os.path.join(BASE_DIR, "..", "assets", "tokenizer"))
-TOKENIZER_DIR = "onnx-community/embeddinggemma-300m-ONNX"
+        model_dir = os.path.dirname(self.model_path)
 
-# MODEL_DIR = os.path.abspath(
-#     os.path.join(BASE_DIR, "..", "assets", "onnx", "model.onnx")
-# )
+        print("🔵 Creating inference session…")
+        self.session = ort.InferenceSession(
+            self.model_path,
+            providers=["CPUExecutionProvider"],
+        )
 
+        self.input_names = [i.name for i in self.session.get_inputs()]
+        self.output_names = [o.name for o in self.session.get_outputs()]
 
-class EmbeddingModel:
-    def __init__(self):
-        # print(TOKENIZER_DIR)
-        self.tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_DIR)
-
-        # sess_options = ort.SessionOptions()
-        # providers = ["CPUExecutionProvider"]
-        #
-        # self.session = ort.InferenceSession(
-        #     MODEL_DIR, sess_options, providers=providers
-        # )
-        #
-        # self.input_names = [inp.name for inp in self.session.get_inputs()]
-        # self.output_names = [out.name for out in self.session.get_outputs()]
-
-    # def _run_sync(
-    #     self, input_ids: np.ndarray, attention_mask: np.ndarray
-    # ) -> List[float]:
-    #     inputs = {}
-    #
-    #     if "input_ids" in self.input_names:
-    #         inputs["input_ids"] = input_ids
-    #     else:
-    #         inputs[self.input_names[0]] = input_ids
-    #
-    #     if "attention_mask" in self.input_names:
-    #         inputs["attention_mask"] = attention_mask
-    #     elif len(self.input_names) > 1:
-    #         inputs[self.input_names[1]] = attention_mask
-    #
-    #     outputs = self.session.run(self.output_names, inputs)
-    #     emb = outputs[0]
-    #
-    #     if emb.ndim == 3:
-    #         emb_vector = emb.mean(axis=1)[0]
-    #     elif emb.ndim == 2:
-    #         emb_vector = emb[0]
-    #     else:
-    #         emb_vector = np.asarray(emb).flatten()
-    #
-    #     return emb_vector.astype(float).tolist()
-
-    async def embed_text(self, text: str, max_length: int = 512) -> List[float]:
+    async def embed_text(self, text: str, max_length=512) -> List[float]:
 
         encoded = self.tokenizer(
             text,
-            return_tensors="np",
             truncation=True,
-            padding="longest",
+            padding=True,
             max_length=max_length,
+            return_tensors="np",
         )
 
         input_ids = encoded["input_ids"].astype(np.int64)
-        attention_mask = encoded.get("attention_mask", np.ones_like(input_ids)).astype(
-            np.int64
+        attention_mask = encoded["attention_mask"].astype(np.int64)
+
+        outputs = self.session.run(
+            self.output_names,
+            {
+                self.input_names[0]: input_ids,
+                self.input_names[1]: attention_mask,
+            },
         )
+        last_hidden = outputs[0]
 
-        # loop = asyncio.get_event_loop()
-        # vector = await loop.run_in_executor(
-        #     None, self._run_sync, input_ids, attention_mask
-        # )
-        # return vector
-        return input_ids.flatten().tolist()
-
+        mask = attention_mask[..., None]
+        pooled = (last_hidden * mask).sum(axis=1) / mask.sum(axis=1)
 
-def cleanup(self):
-    if self.session:
-        self.session = None
-        print("ONNX runtime session closed.")
+        vec = pooled[0]
 
+        norm = np.linalg.norm(vec)
+        if norm > 0:
+            vec = vec / norm
 
-embedding_model = EmbeddingModel()
+        return vec.tolist()
 
 
-async def test_tokenizer():
-    text = "What does the company telll about moonlighting"
-    tokens = await embedding_model.embed_text(text)
-    print("Tokenized text:", tokens)
+embedding_model = EmbeddingModel()

src/chatbot/router.py

@@ -8,7 +8,8 @@ from sqlalchemy import text
 from sqlmodel.ext.asyncio.session import AsyncSession
 
 from src.core.database import get_async_session
-
+from .schemas import ManualTextRequest
+from .service import store_manual_text
 from .embedding import embedding_model
 from .schemas import (
     SemanticSearchRequest,
@@ -21,42 +22,6 @@ from .service import process_pdf_and_store
 
 router = APIRouter(prefix="/chatbot", tags=["chatbot"])
 
-
-# before hitting this endpoint make sure the model.data & model.onnx_data is available on the asset/onnx folder
-@router.post("/upload-pdf", response_model=UploadKBResponse)
-async def upload_pdf(
-    file: UploadFile = File(...),
-    name: str = Form(...),
-    description: Optional[str] = Form(None),
-    session: AsyncSession = Depends(get_async_session),
-):
-    if not file.filename.endswith(".pdf"):
-        raise HTTPException(
-            status_code=400, detail="Only PDF files are supported for now."
-        )
-
-    tmp_dir = tempfile.mkdtemp()
-    tmp_path = os.path.join(tmp_dir, file.filename)
-    try:
-        with open(tmp_path, "wb") as out_f:
-            shutil.copyfileobj(file.file, out_f)
-
-        with open(tmp_path, "rb") as fobj:
-            result = await process_pdf_and_store(fobj, name, description, session)
-
-        return UploadKBResponse(
-            kb_id=result["kb_id"],
-            name=result["name"],
-            chunks_stored=result["chunks_stored"],
-        )
-    finally:
-        try:
-            os.remove(tmp_path)
-            os.rmdir(tmp_dir)
-        except Exception:
-            pass
-
-
 @router.post("/tokenize", response_model=TokenizeResponse)
 async def tokenize_text(payload: TokenizeRequest):
     try:
@@ -88,14 +53,14 @@ async def semantic_search(
     q_vector = payload.embedding
     top_k = payload.top_k or 3
 
-    # Convert Python list → pgvector string format
     q_vector_str = "[" + ",".join(str(x) for x in q_vector) + "]"
 
     sql = text(
         """
-        SELECT id, kb_id, chunk_text, embedding <=> :query_vec AS score
+        SELECT id, kb_id, chunk_text,
+           embedding <#> :query_vec AS score
         FROM knowledge_chunk
-        ORDER BY embedding <=> :query_vec
+        ORDER BY embedding <#> :query_vec ASC
         LIMIT :top_k
         """
     )
@@ -104,7 +69,7 @@ async def semantic_search(
         sql, {"query_vec": q_vector_str, "top_k": top_k}
     )
     rows = result.fetchall()
-
+    
     return [
         SemanticSearchResult(
             chunk_id=str(r.id),
@@ -115,3 +80,79 @@ async def semantic_search(
         for r in rows
     ]
 
+# before hitting this endpoint make sure the model.data & model.onnx_data is available on the asset/onnx folder
+# @router.post("/upload-pdf", response_model=UploadKBResponse)
+# async def upload_pdf(
+#     file: UploadFile = File(...),
+#     name: str = Form(...),
+#     description: Optional[str] = Form(None),
+#     session: AsyncSession = Depends(get_async_session),
+# ):
+#     if not file.filename.endswith(".pdf"):
+#         raise HTTPException(
+#             status_code=400, detail="Only PDF files are supported for now."
+#         )
+
+#     tmp_dir = tempfile.mkdtemp()
+#     tmp_path = os.path.join(tmp_dir, file.filename)
+#     try:
+#         with open(tmp_path, "wb") as out_f:
+#             shutil.copyfileobj(file.file, out_f)
+
+#         with open(tmp_path, "rb") as fobj:
+#             result = await process_pdf_and_store(fobj, name, description, session)
+
+#         return UploadKBResponse(
+#             kb_id=result["kb_id"],
+#             name=result["name"],
+#             chunks_stored=result["chunks_stored"],
+#         )
+#     finally:
+#         try:
+#             os.remove(tmp_path)
+#             os.rmdir(tmp_dir)
+#         except Exception:
+#             pass
+
+# @router.post("/manual-add-chunk")
+# async def manual_add_chunk(
+#     payload: ManualTextRequest,
+#     session: AsyncSession = Depends(get_async_session)
+# ):
+#     return await store_manual_text(
+#         kb_id=payload.kb_id,
+#         text=payload.text,
+#         session=session
+#     )
+
+# @router.post("/test-semantic", response_model=list[SemanticSearchResult])
+# async def test_semantic(
+#     query: str,
+#     top_k: int = 3,
+#     session: AsyncSession = Depends(get_async_session)
+# ):
+
+#     embedding = await embedding_model.embed_text(query)
+
+#     q_vec = "[" + ",".join(map(str, embedding)) + "]"
+
+#     sql = text("""
+#         SELECT id, kb_id, chunk_text,
+#                embedding <#> :vec AS score
+#         FROM knowledge_chunk
+#         ORDER BY embedding <#> :vec ASC
+#         LIMIT :k
+#     """)
+
+#     result = await session.execute(sql, {"vec": q_vec, "k": top_k})
+#     rows = result.fetchall()
+
+#     return [
+#         SemanticSearchResult(
+#             chunk_id=str(r.id),
+#             kb_id=str(r.kb_id),
+#             text=r.chunk_text,
+#             score=float(r.score),
+#         )
+#         for r in rows
+#     ]

src/chatbot/schemas.py

@@ -34,3 +34,7 @@ class SemanticSearchResult(BaseModel):
     kb_id: str
     text: str
     score: float
+
+class ManualTextRequest(BaseModel):
+    kb_id: uuid.UUID
+    text: str

src/chatbot/service.py

@@ -1,7 +1,7 @@
 import os
-
+from uuid import UUID
 from sqlmodel.ext.asyncio.session import AsyncSession
-
+from sqlmodel import select
 from .embedding import embedding_model
 from .models import KnowledgeBase, KnowledgeChunk
 from .utils import (
@@ -43,3 +43,29 @@ async def process_pdf_and_store(
     await session.commit()
 
     return {"kb_id": kb.id, "name": kb_name, "chunks_stored": len(chunk_objs)}
+
+async def store_manual_text(kb_id: UUID, text: str, session: AsyncSession):
+    embedding = await embedding_model.embed_text(text)
+
+    result = await session.execute(
+        select(KnowledgeChunk).where(KnowledgeChunk.kb_id == kb_id)
+    )
+    existing = result.scalars().all()
+    next_index = len(existing)
+
+    new_chunk = KnowledgeChunk(
+        kb_id=kb_id,
+        chunk_index=next_index,
+        chunk_text=text,
+        embedding=embedding
+    )
+
+    session.add(new_chunk)
+    await session.commit()
+
+    return {
+        "kb_id": kb_id,
+        "chunk_index": next_index,
+        "status": "stored",
+        "text": text
+    }

src/main.py

@@ -13,7 +13,7 @@ app = FastAPI(title="Yuvabe App API")
 
 app.include_router(home_router, prefix="/home", tags=["Home"])
 
-init_db()
+# init_db()
 
 app.include_router(auth_router)