updated example

franciscojavierarceo · franciscojavierarceo · commit b7c3426e91e0 · 2025-01-10T16:59:23.000-05:00
Signed-off-by: Francisco Javier Arceo &lt;farceo@redhat.com&gt;
diff --git a/examples/rag/feature_repo/example_repo.py b/examples/rag/feature_repo/example_repo.py
@@ -29,6 +29,7 @@
             vector_index=True,
             vector_search_metric="L2",
         ),
+        Field(name="state", dtype=String),
         Field(name="sentence_chunks", dtype=String),
         Field(name="wiki_summary", dtype=String),
     ],
diff --git a/examples/rag/feature_repo/feature_store.yaml b/examples/rag/feature_repo/feature_store.yaml
@@ -3,8 +3,10 @@ provider: local
 registry: data/registry.db
 online_store:
   type: milvus
+  path: data/online_store.db
   vector_enabled: true
   embedding_dim: 384
+  index_type: "IVF_FLAT"
 
 
 offline_store:
diff --git a/examples/rag/feature_repo/test_workflow.py b/examples/rag/feature_repo/test_workflow.py
@@ -5,6 +5,7 @@
 from pymilvus import MilvusClient, DataType, FieldSchema
 from transformers import AutoTokenizer, AutoModel
 from example_repo import city_embeddings_feature_view, item
+
 TOKENIZER = "sentence-transformers/all-MiniLM-L6-v2"
 MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 
@@ -42,38 +43,10 @@ def run_demo():
     print('\ndata=')
     print(df.head().T)
 
-    store.apply([city_embeddings_feature_view, item])
+    # store.apply([city_embeddings_feature_view, item])
     store.write_to_online_store("city_embeddings", df)
 
-    client = MilvusClient(uir="http://localhost:19530", token="username:password")
-    fields = [
-        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
-        FieldSchema(name='state', dtype=DataType.STRING, description="State"),
-        FieldSchema(name='wiki_summary', dtype=DataType.STRING, description="State"),
-        FieldSchema(name='sentence_chunks', dtype=DataType.STRING, description="Sentence Chunks"),
-        FieldSchema(name="item_id", dtype=DataType.INT64, default_value=0, description="Item"),
-        FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=embedding_length, description="vector")
-    ]
-    cols = [f.name for f in fields]
-    client.insert(
-        collection_name="demo_collection",
-        data=df[cols].to_dict(orient="records"),
-        schema=fields,
-    )
-    print('\n')
-    print('collections', client.list_collections())
-    print('query results =', client.query(
-        collection_name="rag_city_embeddings",
-        filter="item_id == 0",
-        # output_fields=['city_embeddings', 'item_id', 'city_name'],
-    ))
-    print('query results2 =', client.query(
-        collection_name="rag_city_embeddings",
-        filter="item_id >= 0",
-        output_fields=["count(*)"]
-        # output_fields=['city_embeddings', 'item_id', 'city_name'],
-    ))
-    question = "the most populous city in the U.S. state of Texas?"
+    question = "the most populous city in the state of New York is New York"
     tokenizer = AutoTokenizer.from_pretrained(TOKENIZER)
     model = AutoModel.from_pretrained(MODEL)
     query_embedding = run_model(question, tokenizer, model)
@@ -82,12 +55,17 @@ def run_demo():
     # Retrieve top k documents
     features = store.retrieve_online_documents(
         feature=None,
-        features=["city_embeddings:vector", "city_embeddings:item_id", "city_embeddings:state"],
+        features=[
+            "city_embeddings:vector",
+            "city_embeddings:item_id",
+            "city_embeddings:state",
+        ],
         query=query,
-        top_k=3
+        top_k=3,
     )
-    print("features", features.to_df())
-
+    print("features =")
+    print(features.to_df())
+    # store.teardown()
 
 if __name__ == "__main__":
     run_demo()
