updating test_workflow

Signed-off-by: Francisco Javier Arceo <[email protected]>

examples/rag/feature_repo/test_workflow.py

@@ -2,7 +2,6 @@
 import torch
 import torch.nn.functional as F
 from feast import FeatureStore
-from pymilvus import MilvusClient, DataType, FieldSchema
 from transformers import AutoTokenizer, AutoModel
 from example_repo import city_embeddings_feature_view, item
 
@@ -40,11 +39,14 @@ def run_demo():
     embedding_length = len(df['vector'][0])
     print(f'embedding length = {embedding_length}')
 
+    store.apply([city_embeddings_feature_view, item])
+    fields = [
+        f.name for f in city_embeddings_feature_view.features
+    ] + city_embeddings_feature_view.entities + [city_embeddings_feature_view.batch_source.timestamp_field]
     print('\ndata=')
-    print(df.head().T)
+    print(df[fields].head().T)
+    store.write_to_online_store("city_embeddings", df[fields][0:3])
 
-    # store.apply([city_embeddings_feature_view, item])
-    store.write_to_online_store("city_embeddings", df)
 
     question = "the most populous city in the state of New York is New York"
     tokenizer = AutoTokenizer.from_pretrained(TOKENIZER)
@@ -59,13 +61,15 @@ def run_demo():
             "city_embeddings:vector",
             "city_embeddings:item_id",
             "city_embeddings:state",
+            "city_embeddings:sentence_chunks",
+            "city_embeddings:wiki_summary",
         ],
         query=query,
         top_k=3,
     )
     print("features =")
     print(features.to_df())
-    # store.teardown()
+    store.teardown()
 
 if __name__ == "__main__":
     run_demo()