from schema_filter import filter_func, SchemaItemClassifierInference

# 在eval模式下，sql不用提供
data = {
  "text": "Name movie titles released in year 1945. Sort the listing by the descending order of movie popularity.",
  "sql": "",
  "schema": {
    "schema_items": [
      {
        "table_name": "lists",
        "table_comment": "",
        "column_names": [
          "user_id",
          "list_id",
          "list_title",
          "list_movie_number",
          "list_update_timestamp_utc",
          "list_creation_timestamp_utc",
          "list_followers",
          "list_url",
          "list_comments",
          "list_description",
          "list_cover_image_url",
          "list_first_image_url",
          "list_second_image_url",
          "list_third_image_url"
        ],
        "column_comments": [
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          ""
        ]
      },
      {
        "table_name": "movies",
        "table_comment": "",
        "column_names": [
          "movie_id",
          "movie_title",
          "movie_release_year",
          "movie_url",
          "movie_title_language",
          "movie_popularity",
          "movie_image_url",
          "director_id",
          "director_name",
          "director_url"
        ],
        "column_comments": [
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          ""
        ]
      },
      {
        "table_name": "ratings_users",
        "table_comment": "",
        "column_names": [
          "user_id",
          "rating_date_utc",
          "user_trialist",
          "user_subscriber",
          "user_avatar_image_url",
          "user_cover_image_url",
          "user_eligible_for_trial",
          "user_has_payment_method"
        ],
        "column_comments": [
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          ""
        ]
      },
      {
        "table_name": "lists_users",
        "table_comment": "",
        "column_names": [
          "user_id",
          "list_id",
          "list_update_date_utc",
          "list_creation_date_utc",
          "user_trialist",
          "user_subscriber",
          "user_avatar_image_url",
          "user_cover_image_url",
          "user_eligible_for_trial",
          "user_has_payment_method"
        ],
        "column_comments": [
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          ""
        ]
      },
      {
        "table_name": "ratings",
        "table_comment": "",
        "column_names": [
          "movie_id",
          "rating_id",
          "rating_url",
          "rating_score",
          "rating_timestamp_utc",
          "critic",
          "critic_likes",
          "critic_comments",
          "user_id",
          "user_trialist",
          "user_subscriber",
          "user_eligible_for_trial",
          "user_has_payment_method"
        ],
        "column_comments": [
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          "",
          ""
        ]
      }
    ]
  }
}

dataset = [data]

# 最多保留数据库中的7张表
num_top_k_tables = 7
# 对于每张保留的表，最多保留其中20个列，所以输入的prompt中最多有7*10=70个列
num_top_k_columns = 10

# 加载分类器模型
sic = SchemaItemClassifierInference("sic_merged")

# 对于测试数据，我们需要加载训练好的分类器，根据用户问题对表和列打分
dataset = filter_func(
    dataset = dataset, 
    dataset_type = "eval",
    sic = sic,
    num_top_k_tables = num_top_k_tables,
    num_top_k_columns = num_top_k_columns
)

print(dataset)