Skip to content

Commit

Permalink
fix: set model name
Browse files Browse the repository at this point in the history
  • Loading branch information
jamescalam committed Nov 19, 2024
1 parent cf3f771 commit b63f077
Showing 1 changed file with 7 additions and 5 deletions.
12 changes: 7 additions & 5 deletions tests/unit/test_chunkers.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,16 @@
StatisticalChunker,
)

ENCODER_NAME = "text-embedding-3-small"


def test_consecutive_sim_splitter():
# Create a Mock object for the encoder
mock_encoder = Mock()
mock_encoder.return_value = np.array([[1, 0], [1, 0.1], [0, 1]])

encoder = OpenAIEncoder(
name="",
name=ENCODER_NAME,
openai_api_key="a",
)
# Instantiate the ConsecutiveSimSplitter with the mock encoder
Expand Down Expand Up @@ -55,7 +57,7 @@ async def async_return(*args, **kwargs):
mock_encoder.acall.side_effect = async_return

encoder = OpenAIEncoder(
name="",
name=ENCODER_NAME,
openai_api_key="a",
)
# Instantiate the ConsecutiveSimSplitter with the mock encoder
Expand Down Expand Up @@ -91,7 +93,7 @@ def test_cumulative_sim_splitter():

# Instantiate the CumulativeSimSplitter with the mock encoder
encoder = OpenAIEncoder(
name="",
name=ENCODER_NAME,
openai_api_key="a",
)
splitter = CumulativeChunker(encoder=encoder, score_threshold=0.9)
Expand Down Expand Up @@ -127,7 +129,7 @@ async def test_async_cumulative_sim_splitter():

# Instantiate the CumulativeSimSplitter with the mock encoder
encoder = OpenAIEncoder(
name="",
name=ENCODER_NAME,
openai_api_key="a",
)
splitter = CumulativeChunker(encoder=encoder, score_threshold=0.9)
Expand Down Expand Up @@ -181,7 +183,7 @@ def test_statistical_chunker():
mock_encoder.side_effect = lambda docs: np.array([[1, 0] for _ in docs])

encoder = OpenAIEncoder(
name="",
name=ENCODER_NAME,
openai_api_key="a",
)
# Instantiate the ConsecutiveSimSplitter with the mock encoder
Expand Down

0 comments on commit b63f077

Please sign in to comment.