from typing import Any, Dict, Generator from _pytest.fixtures import SubRequest import pytest import pytest_asyncio import numpy as np # The "args" values in TEST_CONFIGS are not optimized for speed but only check that the inference is working for the different models architectures. TEST_CONFIGS = { "BAAI/bge-large-en-v1.5": { "model_id": "BAAI/bge-large-en-v1.5", "input": "What is Deep Learning?", "expected_output": "[[0.016890164, 0.0028579393, -0.051288426, -0.035222173, 0.014418433, -0.02420237, -0.013903489, -0.02162765, -0.005123693, 0.06467697, 0.013903489, 0.0294548, 0.024305359, -0.035016194, -0.01987684, -0.0042482885, -0.016375221, -0.0060505923, -0.046138987, -0.003128285, -0.027189046, 0.051700383, -0.027601, -0.0042482885, -0.038517814, 0.011843713, 0.0708563, 0.027395023, 0.04840474, 0.03460424, -0.019464884, -0.02595318, 0.007621172, -0.032544464, -0.019567873, -0.012410152, 0.010968308, -0.024511335, -0.06302915, -0.05355418, -0.01462441, 0.0313086, 0.023893403, -0.0939258, -0.05190636, 0.010659342, 0.038723793, 0.0026004673, -0.015963266, -0.0032827682, -0.00865106, 0.030278709, -0.009011521, -0.014315444, 0.01555131, -0.021833627, 0.015345332, 0.028424911, 0.018846951, -0.0042482885, 0.045315076, -0.020494772, 0.018743964, -0.041607477, 0.029866755, 0.002162765, 0.010195892, -0.042019434, -0.020391785, -0.028012956, -0.051288426, -0.005329671, 0.00947497, -0.038311835, -0.01555131, 0.024614325, -0.005896109, 0.042843346, -0.04387323, 0.007878643, 0.016169243, 0.04922865, -0.025541224, 0.0294548, -0.016375221, -0.031514574, 0.04387323, 0.033368375, -0.005844615, -0.028424911, 0.020494772, 0.022760527, 0.0021241442, 0.020185806, -0.019464884, 0.026159158, -0.0156543, 0.02502628, 0.0022142595, 0.023996392, 0.018229019, 0.06138133, -0.016581198, 0.08156714, -0.04922865, 0.0096294535, 0.015860276, -0.005381165, -0.028012956, -0.047168873, -0.014418433, -0.001029888, -0.011895208, -0.016890164, -0.035634127, 0.025541224, -0.008342094, 0.016787175, -0.0066942726, 0.0065912837, 0.023687426, 0.014521422, 0.022760527, 0.019258907, 0.015963266, -0.047786806, -0.019670863, 0.0030124227, -0.010247386, 0.0019181665, -0.0008882785, -0.033986308, -0.011843713, 0.036046084, 0.00368185, -0.042843346, 0.008445082, 0.03975368, -0.021936616, 0.0020726498, 0.01987684, -0.011843713, -0.011895208, 0.058085687, -0.048816696, 0.0045315074, -0.009114509, 0.010092903, -0.009577959, 0.02070075, -0.026983067, 0.009680948, 0.024511335, 0.006102087, -0.07168021, -0.042431388, -0.027601, 0.036664017, -0.003321389, 0.0074666888, 0.0115862405, 0.006385306, -0.032750443, 0.04304932, -0.0020211553, 0.05005256, -0.04037161, -0.04387323, 0.01462441, -0.016890164, 0.02924882, -0.008960026, -0.026365135, -0.01297659, -0.011792218, 0.057261776, 0.05005256, -0.021215694, 0.025438236, -0.012925096, -0.010401869, -0.033368375, -0.0138005, 0.0040423106, 0.0020082816, -0.014830388, -0.022245582, 0.016890164, -0.06673675, 0.026365135, 0.014212456, 0.016684188, -0.041195523, 0.011740725, -0.016478209, 0.020288795, -0.03439826, 0.0043770242, 0.0087540485, -0.042225413, 0.011946701, 0.04655094, -0.0034243779, -0.0032698947, -0.013182567, 0.070444345, 0.016684188, 0.055613957, -0.06426501, 0.019155918, 0.02070075, 0.00782715, -0.023790415, 0.009114509, 0.044285186, -0.042843346, -0.067972615, -0.0011843713, 0.0065397895, -0.02337846, 0.019773852, 0.02152466, 0.038723793, 0.007878643, -0.015757287, -0.02677709, -0.009011521, 0.02348145, -0.028424911, 0.017508097, 0.013594523, 0.07827149, -0.011019803, 0.018846951, 0.016478209, 0.023069493, 0.010350375, 0.04078357, 0.019464884, 0.035016194, -0.0136975115, -0.020185806, 0.02502628, 0.028630888, -0.00066299044, -0.024511335, 0.014315444, 0.003553114, 0.0049949572, 0.013285556, -0.030072732, 0.012616129, 0.026159158, -0.0017121889, -0.032544464, 0.04572703, 0.044491164, 0.030072732, -0.032956418, -0.042637367, -0.019670863, 0.06220524, -0.011998196, -0.01740511, 0.01905293, -0.007672666, 0.016787175, 0.0031154114, -0.028218934, -0.027395023, -0.03707597, -0.07909541, -0.063441105, -0.018640975, -0.026365135, -0.008290599, 0.006642778, -0.02595318, 0.01168923, -0.006359559, 0.007672666, -0.006462548, -0.012101185, -0.001383912, 0.040989544, 0.019155918, -0.04037161, 0.024099382, -0.036046084, -0.004840474, -0.010659342, -0.009217498, -0.01987684, -0.0035016194, 0.013903489, -0.027189046, 0.02924882, -0.02502628, -0.0294548, -0.020597761, 0.014109467, -0.039341725, -0.032132506, 0.04655094, -0.02595318, 0.024820304, 0.032956418, -0.0020469027, 0.028836867, 0.011946701, -0.053966135, 0.029042844, 0.018434996, 0.013388545, -0.03110262, 0.0313086, 0.023687426, -0.028218934, 0.004222541, -0.057261776, -0.040989544, -0.0011972449, 0.007621172, -0.07332803, 0.010350375, 0.046138987, 0.0043770242, -0.08568669, 0.021421673, -0.03460424, -0.015757287, -0.011174286, 0.040165637, 0.029866755, 0.035840105, 0.019258907, -0.03378033, 0.0034243779, -0.02348145, 0.027806979, 0.03481022, -0.04057759, 0.024923291, 0.03625206, 0.0052781766, 0.01730212, -0.020803738, -0.01905293, 0.034192283, 0.028012956, 0.03378033, 0.0294548, 0.009680948, -0.010607847, 0.01905293, -0.027806979, -0.03975368, 0.052318316, 0.021730639, -0.0030896643, 0.037487928, 0.04057759, 0.00034115041, 0.014109467, -0.010968308, 0.0021885121, -0.021730639, 0.015860276, 0.047992785, -0.04037161, 0.026983067, 0.008496577, -0.038929768, 0.035222173, -0.052318316, -0.028218934, 0.014109467, -0.03728195, 0.021215694, -0.07003239, -0.0065397895, 0.026983067, -0.011740725, 0.026056169, 0.024099382, 0.04840474, -0.022966504, 0.007569677, -0.022142593, 0.012101185, 0.023069493, -0.014418433, -0.058085687, 0.016375221, -0.026983067, -0.017508097, 0.05973351, 0.0023816163, 0.0313086, -0.018332008, 0.08074322, -0.012152679, 0.020082818, 0.008290599, -0.0026004673, 0.0035016194, -0.032956418, -0.0059218565, 0.03439826, -0.003257021, -0.00865106, 0.009114509, 0.032132506, -0.042225413, -0.026365135, -0.004480013, 0.0049949572, -0.052112337, -0.021833627, 0.0626172, -0.02502628, 0.011328769, -0.08898233, 0.016066253, 0.02348145, -0.060969375, 0.03378033, -0.028836867, 0.055613957, 0.025644213, -0.028012956, -0.038517814, 0.0023043745, -0.015963266, -0.03810586, 0.0045057605, 0.053966135, -0.03481022, 0.040165637, -0.038517814, -0.020906728, -0.0023043745, 0.06673675, -0.0028965601, -0.019155918, -0.023172481, -0.0136975115, 0.038311835, 0.011071296, -0.009063015, -0.009166004, -0.040165637, 0.04902267, -0.028424911, -0.03110262, 0.023584437, -0.010659342, -0.030690664, 0.011122791, 0.013903489, 0.030072732, -0.0033471363, 0.05602591, -0.046138987, -0.043667253, 0.044491164, -0.022554548, -0.0047632325, 0.042637367, 0.010813825, -0.025644213, -0.00625657, 0.057673734, -0.051082447, 0.052112337, -0.0156543, 0.024408348, -0.011019803, -0.029866755, -0.02152466, -0.04675692, 0.04407921, -7.724161e-05, -0.002484605, -0.019361896, -0.05252429, -0.02348145, 0.011483252, -0.035222173, -0.0021498913, -0.032956418, 0.012616129, 0.027189046, 0.010453364, 0.01076233, 0.056849822, -0.005123693, 0.027806979, 0.035222173, 0.012719118, 0.0294548, -0.033368375, -0.0589096, -0.0030896643, -0.032132506, 0.007003239, -0.033368375, 0.022863515, -0.008496577, 0.024408348, -0.020906728, 0.0023429955, 0.006127834, 0.018846951, 0.022863515, -0.010968308, -0.022863515, 0.0042740353, -0.0029866754, 0.07168021, 0.032338485, -0.04655094, -0.07332803, 0.044491164, -0.046344962, 0.009680948, 0.0059218565, -0.04922865, 0.014521422, -0.040989544, 0.05355418, -0.049846582, -0.015139354, -0.032750443, -0.0016864417, -0.007621172, 0.01555131, -0.01730212, -0.0051494404, 0.025850192, -0.044285186, 0.011998196, -0.001029888, 0.008702555, -0.06385306, 0.0047632325, -0.02924882, 0.05005256, 0.0068487558, -0.038311835, 0.006642778, -0.020803738, 0.033986308, 0.020288795, -0.06879652, -0.02348145, 0.01555131, -0.025232257, 0.055202, 0.03460424, -0.01987684, 0.028218934, -0.013028084, 0.0040165633, -0.03378033, 0.016478209, -0.0028836867, -0.04037161, 0.04057759, -0.059321553, 0.010659342, 0.004737485, 0.015448322, 0.04758083, 0.008290599, -0.042225413, 0.00496921, -0.026056169, -0.074975856, -0.025747202, -0.03481022, -0.00046023124, -0.0011714977, -0.012770612, 0.039135747, -0.008908532, 0.045521054, 0.022142593, 0.034192283, -0.070444345, -0.007621172, -0.025232257, 0.036664017, -0.027806979, 0.0058703623, -0.03686999, -0.0021112706, -0.03357435, -0.03686999, -0.043461278, -0.01297659, 0.020803738, 0.047374852, -0.007415194, 0.009886926, -0.019464884, -0.016478209, -0.03378033, -0.020803738, 0.024511335, 0.0138005, 0.03975368, 0.019567873, 0.0033471363, 0.070444345, -0.023584437, -0.013182567, 0.017920053, 0.07209217, -0.031720553, -0.01812603, -0.014315444, 0.031514574, -0.026983067, -0.0313086, 0.049640607, -0.017508097, 0.043461278, -0.028630888, 0.021936616, -0.06426501, 0.028012956, -0.004171047, 0.0149333775, -0.0046344963, 0.0061793285, 0.04655094, 0.023172481, -0.04675692, 0.029042844, 0.042019434, -0.003578861, 0.010710836, -0.068384565, -0.06138133, -0.027806979, -0.048198763, 0.034192283, -0.054378092, -0.019979829, 0.06138133, 0.004454266, 0.04593301, 0.011277274, -0.015139354, 0.0055099013, -0.014315444, -0.054790046, 0.028424911, 0.0049177157, 0.020288795, 0.0136975115, -0.030690664, -0.015448322, 0.0087540485, 0.03625206, 0.007878643, -0.04078357, -0.008908532, -0.050876472, -0.037487928, -0.025541224, -0.011019803, -0.015963266, 0.009114509, 0.021318683, -0.004119552, 0.032132506, 0.010041409, -0.04943463, -0.026365135, -0.045521054, 0.041195523, -0.044697143, 0.0033728834, 0.010298881, -0.017817063, 0.0023043745, 0.025850192, 0.02337846, 0.01905293, 0.0033986308, 0.0035016194, -0.015242344, 0.038723793, -0.035840105, 0.022554548, -0.044697143, 0.019773852, -0.0013195442, 0.01905293, 0.020494772, -0.054790046, -0.013594523, -0.035840105, -0.01812603, -0.03481022, 0.011277274, -0.012667623, 0.0058703623, 0.0149333775, -0.028218934, -0.00020517301, -0.024099382, -0.038723793, -0.014315444, -0.038929768, 0.013491534, -0.013182567, 0.04758083, 0.015963266, 0.02162765, 0.05973351, 0.026571112, 0.03357435, -0.02924882, -0.039135747, 0.022451561, -0.0054584066, -0.030690664, -0.018846951, -0.025644213, 0.01555131, -0.03439826, 0.012564635, 0.010659342, -0.008908532, -0.021833627, 0.014006478, -0.028630888, 0.024820304, 0.0014675906, 0.003810586, -0.0034243779, 0.015860276, -0.03789988, -0.016478209, -0.022039605, -0.009011521, 0.011328769, -0.03378033, -0.013388545, -0.020288795, -0.022554548, 0.023790415, 0.033986308, 0.012564635, -0.01894994, 0.003450125, 0.03707597, -0.007569677, -0.020906728, -0.011328769, -0.03460424, 0.008702555, -0.03728195, -0.03707597, 0.02924882, 0.033986308, 0.0589096, -0.0031669058, -0.018332008, 0.017817063, -0.010401869, 0.055613957, -0.019464884, 0.060969375, -0.0042482885, 0.004711738, 0.02502628, 0.016169243, 2.3132252e-05, 0.09145406, 0.030278709, -0.004737485, -0.017817063, -0.026262145, 0.014521422, -0.0469629, 0.024511335, -0.01987684, 0.038311835, 0.018332008, 0.016272232, 0.037693903, 0.054790046, 0.0432553, 0.0038620804, -0.02162765, 0.0020726498, 0.04037161, -0.010710836, 0.044697143, 0.041607477, 0.030690664, 0.010453364, 0.016478209, -0.034192283, -0.00012712681, -0.00690025, -0.022348572, -0.034192283, -0.0009526465, -0.03460424, -0.00031862163, -0.018434996, 0.035016194, 0.01462441, -0.033162396, -0.00690025, 0.035016194, 0.0589096, 0.028012956, 0.03707597, -0.0043512774, -0.013491534, 0.0041452996, 0.005741626, 0.032956418, 0.019567873, -0.001486901, -0.027189046, 0.034192283, -0.021009717, 0.030690664, -0.0036561026, -0.011792218, 0.00432553, 0.057261776, 0.0067972615, -0.04304932, 0.0051751877, -0.0032055266, -0.02070075, -0.034192283, 0.010092903, -0.028424911, 0.00084322086, 0.021215694, -0.022554548, -0.019361896, -0.008805543, -0.012410152, 0.010504859, 0.017714076, 0.021318683, 0.008290599, 0.07373998, 0.053142224, 0.021833627, 0.0046344963, 0.013079579, 0.01812603, 0.0059733507, 0.0070547336, -0.0065912837, -0.022451561, -0.05252429, -0.046344962, 0.018537985, 0.014830388, -0.005896109, 0.007157722, 0.0030252961, 0.059321553, -0.032750443, -0.033986308, 0.03192653, -0.0064110532, 0.009732442, 0.0149333775, -0.0469629, -0.033162396, -0.000102586506, -0.05973351, 0.09886926, -0.013388545, -0.04057759, 0.0022400066, -0.03728195, -0.01462441, 0.025335247, 0.04840474, -0.027806979, 0.03975368, 0.042637367, -0.011328769, 0.00496921, -0.016581198, -0.027395023, 0.030690664, 0.027189046, 0.0070547336, 0.03810586, 0.020906728, 0.008908532, 0.015036366, 0.032338485, -0.006385306, 0.022863515, -0.042431388, -0.016169243, 0.021009717, -0.01894994, -0.03481022, 0.07415194, 0.060145464, -0.02348145, 0.002677709, -0.060557418, 0.0149333775, -0.04675692, 0.011431757, -0.038929768, 0.028836867, -0.041195523, -0.04840474, 0.0044027716, -0.031720553, 0.1812603, 0.053966135, 0.043667253, 0.004737485, 0.07250412, 0.023790415, 0.023584437, 0.034192283, 0.023584437, -0.009371982, 0.08403887, -0.025850192, 0.020597761, -0.011534747, 0.014830388, 0.018434996, 0.0012616129, 0.032750443, -0.03707597, -0.012564635, -0.052112337, -0.0010170145, -0.020288795, -0.042431388, -0.0115862405, -0.0035273668, -0.01462441, -0.039547704, 0.029042844, -0.015963266, -0.005381165, 0.015448322, 0.035840105, -0.012770612, 0.016066253, 0.042843346, 0.039135747, -0.038311835, -0.015448322, 0.045521054, -0.011637735, -0.07703563, 0.019155918, 0.0059218565, -0.01297659, 0.036458038, -0.0149333775, -0.021936616, -0.017920053, -0.047786806, 0.035016194, -0.06632479, -0.018640975, -0.042637367, -0.0012809233, -0.006359559, -0.03728195, -0.01894994, 0.017508097, 0.01730212, 0.021009717, 0.0087540485, -0.036458038, -0.036458038, 0.03728195, 0.007930138, 0.04758083, 0.007106228, -0.0057931202, -0.039135747, 0.006462548, 0.00690025, -0.036458038, 0.030484688, 0.016375221, 0.0626172, -0.031514574, 0.04758083, -0.0313086, -0.035840105, -0.0136975115, 0.00039264484, 0.007878643, -0.022245582, 0.06879652, -0.0073637, -0.03460424, 0.022760527, -0.05973351, 0.011071296, 0.011534747, 0.018537985, 0.05355418, -0.030896643, 0.04902267]]", "env_config": { "MAX_WARMUP_SEQUENCE_LENGTH": "512", }, "args": [ "--dtype", "bfloat16", ], }, } @pytest.fixture(scope="module", params=TEST_CONFIGS.keys()) def test_config(request: SubRequest) -> Dict[str, Any]: """Fixture that provides model configurations for testing.""" model_name = request.param test_config = TEST_CONFIGS[model_name] test_config["test_name"] = model_name return test_config @pytest.fixture(scope="module") def model_id(test_config: Dict[str, Any]) -> Generator[str, None, None]: yield test_config["model_id"] @pytest.fixture(scope="module") def test_name(test_config: Dict[str, Any]) -> Generator[str, None, None]: yield test_config["test_name"] @pytest.fixture(scope="module") def expected_outputs(test_config: Dict[str, Any]) -> Dict[str, str]: return { "expected_output": test_config["expected_output"], } @pytest.fixture(scope="module") def input(test_config: Dict[str, Any]) -> str: return test_config["input"] @pytest.fixture(scope="function") def tei_service(gaudi_launcher, model_id: str, test_name: str): with gaudi_launcher(model_id, test_name) as tei_service: yield tei_service @pytest_asyncio.fixture(scope="function") async def tei_client(tei_service): await tei_service.health(1000) return tei_service @pytest.mark.asyncio async def test_model_single_request( tei_client, expected_outputs: Dict[str, str], input: str ): print(f"type(tei_client) {tei_client}") response = await tei_client.generate( input, ) response_array = np.array(response) expected_array = np.array(eval(expected_outputs["expected_output"])) if not np.allclose(response_array, expected_array, rtol=1e-5, atol=1e-5): print("\nExpected output:") print(f"{expected_array.tolist()}") print("\nReceived output:") print(f"{response_array.tolist()}") raise AssertionError("Response array does not match expected array within tolerance")