File tree Expand file tree Collapse file tree 1 file changed +0
-21
lines changed
tests/integration_tests/models/llm/api Expand file tree Collapse file tree 1 file changed +0
-21
lines changed Original file line number Diff line number Diff line change 11# pylint: disable=redefined-outer-name
2-
3- import math
42from unittest .mock import AsyncMock , MagicMock , patch
53
64import pytest
@@ -81,25 +79,6 @@ async def test_rpm_tpm_limiter_called(ollama_client: OllamaClient):
8179 )
8280
8381
84- @pytest .mark .asyncio
85- async def test_generate_topk_per_token (ollama_client : OllamaClient ):
86- ollama_client .client .chat .return_value = {
87- "message" : {
88- "logprobs" : {
89- "content" : [
90- {"token" : "hello" , "logprob" : - 0.1 },
91- {"token" : "world" , "logprob" : - 0.2 },
92- ]
93- }
94- }
95- }
96- tokens = await ollama_client .generate_topk_per_token ("test" )
97- assert len (tokens ) == 2
98- assert tokens [0 ].text == "hello"
99- assert math .isclose (tokens [0 ].prob , math .exp (- 0.1 ))
100- assert tokens [1 ].text == "world"
101-
102-
10382def test_import_error_when_ollama_missing ():
10483 with patch .dict ("sys.modules" , {"ollama" : None }):
10584 with pytest .raises (ImportError , match = "Ollama SDK is not installed" ):
You can’t perform that action at this time.
0 commit comments