checking in progress

Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
franciscojavierarceo · Mar 21, 2025 · Mar 1, 2025 · Mar 2, 2025 · Mar 2, 2025 · Mar 2, 2025
commit fb1bce4872695d35a68a7c7b44f886fc83523e98
diff --git a/examples/rag-docling/feature_repo/example_repo.py b/examples/rag-docling/feature_repo/example_repo.py
@@ -6,7 +6,7 @@
     FileSource,
 )
 from feast.data_format import ParquetFormat
-from feast.types import Float64, Array, String, ValueType
+from feast.types import Float64, Array, String, ValueType, Bytes
 from feast import Entity, RequestSource
 from feast.on_demand_feature_view import on_demand_feature_view
 from sentence_transformers import SentenceTransformer
@@ -74,3 +74,83 @@ def embed_chunk(inputs: Dict[str, Any]) -> Dict[str, List[float]]:
     source=source,
     ttl=timedelta(hours=2),
 )
+
+import hashlib
+
+def generate_chunk_id(file_name: str, raw_chunk_markdown: str="") -> str:
+    """Generate a unique chunk ID based on file_name and raw_chunk_markdown."""
+    unique_string = f"{file_name}-{raw_chunk_markdown}" if raw_chunk_markdown != "" else f"{file_name}"
+    return hashlib.sha256(unique_string.encode()).hexdigest()
+
+from docling.datamodel.base_models import DocumentStream
+# Load tokenizer and embedding model
+EMBED_MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"
+MAX_TOKENS = 64  # Small token limit for demonstration
+
+import io
+from docling.document_converter import DocumentConverter
+from transformers import AutoTokenizer
+from sentence_transformers import SentenceTransformer
+from docling.chunking import HybridChunker
+
+tokenizer = AutoTokenizer.from_pretrained(EMBED_MODEL_ID)
+embedding_model = SentenceTransformer(EMBED_MODEL_ID)
+chunker = HybridChunker(tokenizer=tokenizer, max_tokens=MAX_TOKENS, merge_peers=True)
+
+input_request_pdf = RequestSource(
+    name="pdf_request_source",
+    schema=[
+        Field(name="pdf_bytes", dtype=Bytes),
+        Field(name="file_name", dtype=String),
+    ],
+)
+
+import pandas as pd
+@on_demand_feature_view(
+    sources=[input_request_pdf],
+    schema=[
+        Field(name="document_id", dtype=String, vector_index=False),
+        Field(name="chunk_id", dtype=Array(String), vector_index=False),
+        Field(name="chunk_text", dtype=Array(String), vector_index=False),
+        Field(name="vector", dtype=Array(Float64), vector_index=False),
+    ],
+    mode="python",
+    singleton=False,
+    write_to_online_store=False,
+    explode=True,
+)
+def docling_transform_docs(inputs: Dict[str, Any]):
+    try:
+        buf = io.BytesIO(
+            inputs['pdf_bytes'],
+        )
+
+        source = DocumentStream(name=inputs['file_name'], stream=buf)
+        converter = DocumentConverter()
+        result = converter.convert(source)
+        document_id = generate_chunk_id(inputs["file_name"])
+        output, chunks, embeddings, chunk_ids = [], [], [], []
+        for chunk in chunker.chunk(dl_doc=result.document):
+            raw_chunk = chunker.serialize(chunk=chunk)
+            embedding = embed_chunk(raw_chunk)
+            chunk_id = generate_chunk_id(inputs["file_name"], raw_chunk)
+            # chunk_dict = {
+            #     "document_id": document_id,
+            #     "chunk_id": chunk_id,
+            #     "vector": embedding,
+            #     "chunk_text": raw_chunk,
+            # }
+            chunks.append(raw_chunk)
+            chunk_ids.append(chunk_id)
+            embeddings.append(embedding)
+            # output.append(chunk_dict)
+        return {"document_id": document_id, "chunk_id": chunks, "vector": embedding, "chunk_text": chunks}
+
+
+    except Exception as e:
+        print(e)
+        return {"document_id": "missing", "chunk_id": ["missing"], "vector": [[0.1, 0.2]], "chunk_text": ["missing"]}
+
+    return pd.DataFrame(output)
+
+
@@ -1000,7 +1000,7 @@ def test_sqlite_get_online_documents_v2_search() -> None:
         assert result["distance"] == [-1.8458267450332642, -1.8458267450332642]
 
 
-@pytest.mark.skip(reason="Skipping this test as CI struggles with it")
+# @pytest.mark.skip(reason="Skipping this test as CI struggles with it")
 def test_local_milvus() -> None:
     import random
 

@@ -19,6 +19,7 @@
     RequestSource,
 )
 from feast.driver_test_data import create_driver_hourly_stats_df
+from feast.nlp_test_data import create_document_chunks_df
 from feast.feature_view import DUMMY_ENTITY_FIELD
 from feast.field import Field
 from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig
@@ -831,6 +832,7 @@ def test_stored_writes(self):
             )
 
             driver = Entity(name="driver", join_keys=["driver_id"])
+            word = Entity(name="word", join_keys=["word_id"])
 
             driver_stats_source = FileSource(
                 name="driver_hourly_stats_source",