fix sorting table relevance scores (#515)

Dataherald · Jun 24, 2024 · f3cb505 · f3cb505
1 parent 7fcd9e8
commit f3cb505
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/services/engine/dataherald/sql_generator/dataherald_sqlagent.py b/services/engine/dataherald/sql_generator/dataherald_sqlagent.py
@@ -294,8 +294,8 @@ def _run(  # noqa: PLR0912
         df["similarities"] = df.table_embedding.apply(
             lambda x: self.cosine_similarity(x, question_embedding)
         )
-        df = df.sort_values(by="similarities", ascending=True)
-        df = df.tail(TOP_TABLES)
+        df = df.sort_values(by="similarities", ascending=False)
+        df = df.head(TOP_TABLES)
         most_similar_tables = self.similar_tables_based_on_few_shot_examples(df)
         table_relevance = ""
         for _, row in df.iterrows():