fix entities

defog-ai · Nov 1, 2024 · 1c2e240 · 1c2e240
1 parent a986ba4
commit 1c2e240
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/test_utils_pruning.py b/tests/test_utils_pruning.py
@@ -85,7 +85,7 @@ def test_metadata_diff_coldesc():
 def test_get_md_emb_no_shuffle(test_metadata):
     column_emb, column_csv, column_ner, column_join = test_metadata
     question = "How many flights start from Los Angeles Airport (LAX)?"
-    assert get_entity_types(question) == {"FAC"}
+    assert get_entity_types(question) == {"GPE", "ORG"}
     k = 3
     threshold = 0.0
 
@@ -129,7 +129,7 @@ def test_get_md_emb_no_shuffle(test_metadata):
 def test_get_md_emb_shuffle(test_metadata):
     column_emb, column_csv, column_ner, column_join = test_metadata
     question = "How many flights start from Los Angeles Airport (LAX)?"
-    assert get_entity_types(question) == {"FAC"}
+    assert get_entity_types(question) == {"GPE", "ORG"}
     k = 3
     threshold = 0.0
 
@@ -195,7 +195,7 @@ def test_get_md_emb_sql_emb_empty(test_metadata):
 def test_get_md_emb_coldesc(test_metadata_diff_coldesc):
     column_emb, column_csv, column_ner, column_join = test_metadata_diff_coldesc
     question = "How many flights start from Los Angeles Airport (LAX)?"
-    assert get_entity_types(question) == {"FAC"}
+    assert get_entity_types(question) == {"GPE", "ORG"}
     k = 3
     threshold = 0.0