Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 1 addition & 4 deletions src/functions/aggregate/llm_rerank/implementation.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -109,10 +109,7 @@ nlohmann::json LlmRerank::SlidingWindow(nlohmann::json& tuples) {
auto carry_forward_tuples = nlohmann::json::array();
int start_index = 0;

auto batch_size = static_cast<int>(model.GetModelDetails().batch_size);
if (batch_size == 2048) {
batch_size = std::min<int>(batch_size, num_tuples);
}
auto batch_size = std::min<int>(model.GetModelDetails().batch_size, num_tuples);

@queryproc queryproc Jun 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should make sure that we take the min of 2: (user_batch, default_batch = 16). Using this batch size, we then expand the data section of the prompt while not going beyond the max input context window size.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Alternatively, if the user provided a batch_size, it overrides the default size and then we expand till we reach the max input context window size.


if (batch_size <= 0) {
throw std::runtime_error("Batch size must be greater than zero");
Expand Down
2 changes: 2 additions & 0 deletions src/include/flock/model_manager/repository.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@

namespace flock {

inline constexpr int DEFAULT_BATCH_SIZE = 16;

struct ModelDetails {
std::string provider_name;
std::string model_name;
Expand Down
2 changes: 1 addition & 1 deletion src/model_manager/model.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ void Model::LoadModelDetails(const nlohmann::json& model_json) {
} else if (db_model_args.contains("batch_size")) {
model_details_.batch_size = db_model_args.at("batch_size").get<int>();
} else {
model_details_.batch_size = 2048;
model_details_.batch_size = DEFAULT_BATCH_SIZE;
}
}
}
Expand Down
3 changes: 2 additions & 1 deletion test/unit/model_manager/model_manager_test.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ TEST_F(ModelManagerTest, ModelInitializationMinimal) {
EXPECT_EQ(details.model_name, "gpt-4o-test");
EXPECT_EQ(details.model, "gpt-4o");
EXPECT_EQ(details.provider_name, "openai");
EXPECT_EQ(details.batch_size, DEFAULT_BATCH_SIZE);
});
}

Expand Down Expand Up @@ -127,4 +128,4 @@ TEST_F(ModelManagerTest, GetModelDetails) {
EXPECT_EQ(details.batch_size, 10);
}

}// namespace flock
}// namespace flock
Loading