Spaces:
Runtime error
Runtime error
File size: 2,542 Bytes
9c49811 f30cbcc 9c49811 f30cbcc 9c49811 f30cbcc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import pandas as pd
import pytest
from utils import filter_models, search_table, filter_queries, select_columns, update_table_long_doc
@pytest.fixture
def toy_df():
return pd.DataFrame(
{
"Retrieval Model": [
"bge-m3",
"bge-m3",
"jina-embeddings-v2-base",
"jina-embeddings-v2-base"
],
"Reranking Model": [
"bge-reranker-v2-m3",
"NoReranker",
"bge-reranker-v2-m3",
"NoReranker"
],
"Average ⬆️": [0.6, 0.4, 0.3, 0.2],
"wiki_en": [0.8, 0.7, 0.2, 0.1],
"wiki_zh": [0.4, 0.1, 0.4, 0.3],
"news_en": [0.8, 0.7, 0.2, 0.1],
"news_zh": [0.4, 0.1, 0.4, 0.3],
}
)
@pytest.fixture
def toy_df_long_doc():
return pd.DataFrame(
{
"Retrieval Model": [
"bge-m3",
"bge-m3",
"jina-embeddings-v2-base",
"jina-embeddings-v2-base"
],
"Reranking Model": [
"bge-reranker-v2-m3",
"NoReranker",
"bge-reranker-v2-m3",
"NoReranker"
],
"Average ⬆️": [0.6, 0.4, 0.3, 0.2],
"law_en_lex_files_300k_400k": [0.4, 0.1, 0.4, 0.3],
"law_en_lex_files_400k_500k": [0.8, 0.7, 0.2, 0.1],
"law_en_lex_files_500k_600k": [0.8, 0.7, 0.2, 0.1],
"law_en_lex_files_600k_700k": [0.4, 0.1, 0.4, 0.3],
}
)
def test_filter_models(toy_df):
df_result = filter_models(toy_df, ["bge-reranker-v2-m3", ])
assert len(df_result) == 2
assert df_result.iloc[0]["Reranking Model"] == "bge-reranker-v2-m3"
def test_search_table(toy_df):
df_result = search_table(toy_df, "jina")
assert len(df_result) == 2
assert df_result.iloc[0]["Retrieval Model"] == "jina-embeddings-v2-base"
def test_filter_queries(toy_df):
df_result = filter_queries("jina", toy_df)
assert len(df_result) == 2
assert df_result.iloc[0]["Retrieval Model"] == "jina-embeddings-v2-base"
def test_select_columns(toy_df):
df_result = select_columns(toy_df, ['news',], ['zh',])
assert len(df_result.columns) == 4
assert df_result['Average ⬆️'].equals(df_result['news_zh'])
def test_update_table_long_doc(toy_df_long_doc):
df_result = update_table_long_doc(toy_df_long_doc, ['law',], ['en',], ["bge-reranker-v2-m3", ], "jina")
print(df_result) |