Skip to content
Snippets Groups Projects
Commit 33e6ff36 authored by Steven Holte's avatar Steven Holte Committed by Chromium LUCI CQ
Browse files

Use the min context limit for Answer input token limit.

(cherry picked from commit ab7d0521)

Bug: 372535307
Change-Id: I3c8ff10d896e309c34b997d3d22793ccbd5fbf1b
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/5950645


Commit-Queue: Steven Holte <holte@chromium.org>
Reviewed-by: default avatarJun Zou <junzou@chromium.org>
Auto-Submit: Steven Holte <holte@chromium.org>
Reviewed-by: default avatarZekun Jiang <zekunjiang@google.com>
Reviewed-by: default avatarOrin Jaworski <orinj@chromium.org>
Cr-Original-Commit-Position: refs/heads/main@{#1372250}
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/5960798
Commit-Queue: Zekun Jiang <zekunjiang@google.com>
Auto-Submit: Zekun Jiang <zekunjiang@google.com>
Cr-Commit-Position: refs/branch-heads/6778@{#787}
Cr-Branched-From: b21671ca-refs/heads/main@{#1368529}
parent 74311268
No related branches found
No related tags found
No related merge requests found
......@@ -141,8 +141,7 @@ class MlAnswerer::SessionManager {
base::OnceCallback<void(int)> session_added_cb,
std::vector<ModelInput> inputs) {
HistoryAnswerRequest request;
// TODO(crbug.com/372535307): use actual model limit.
int token_limit = 1024;
int token_limit = session->GetTokenLimits().min_context_tokens;
// Reserve space for preamble text.
int token_count = kPreambleTokenBufferSize;
......
......@@ -44,8 +44,23 @@ class HistoryEmbeddingsMlAnswererTest : public testing::Test {
public:
void SetUp() override {
ml_answerer_ = std::make_unique<MlAnswerer>(&model_executor_);
token_limits_ = {
.min_context_tokens = 1024,
};
ON_CALL(session_1_, GetTokenLimits())
.WillByDefault([&]() -> optimization_guide::TokenLimits& {
return GetTokenLimits();
});
ON_CALL(session_2_, GetTokenLimits())
.WillByDefault([&]() -> optimization_guide::TokenLimits& {
return GetTokenLimits();
});
}
optimization_guide::TokenLimits& GetTokenLimits() { return token_limits_; }
protected:
optimization_guide::StreamingResponse MakeResponse(
const std::string& answer_text,
......@@ -65,6 +80,7 @@ class HistoryEmbeddingsMlAnswererTest : public testing::Test {
testing::NiceMock<MockModelExecutor> model_executor_;
std::unique_ptr<MlAnswerer> ml_answerer_;
testing::NiceMock<optimization_guide::MockSession> session_1_, session_2_;
optimization_guide::TokenLimits token_limits_;
};
TEST_F(HistoryEmbeddingsMlAnswererTest, ComputeAnswerNoSession) {
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment