Skip to content

Commit

Permalink
FIX: Open AI embedding shortening is only available for some models
Browse files Browse the repository at this point in the history
  • Loading branch information
romanrizzi committed Jan 21, 2025
1 parent 3b66fb3 commit a9d5210
Show file tree
Hide file tree
Showing 3 changed files with 32 additions and 7 deletions.
8 changes: 7 additions & 1 deletion app/models/embedding_definition.rb
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,8 @@ def presets
tokenizer_class: "DiscourseAi::Tokenizer::MultilingualE5LargeTokenizer",
provider: HUGGING_FACE,
},
# "text-embedding-3-large" real dimentions are 3072, but we only support up to 2000 in the
# indexes, so we downsample to 2000 via API.
{
preset_id: "text-embedding-3-large",
display_name: "OpenAI's text-embedding-3-large",
Expand Down Expand Up @@ -198,11 +200,15 @@ def hugging_face_client
end

def open_ai_client
model_name = lookup_custom_param("model_name")
can_shorten_dimensions = %w[text-embedding-3-small text-embedding-3-large].include?(model_name)
client_dimensions = can_shorten_dimensions ? dimensions : nil

DiscourseAi::Inference::OpenAiEmbeddings.new(
endpoint_url,
api_key,
lookup_custom_param("model_name"),
dimensions,
client_dimensions,
)
end

Expand Down
28 changes: 25 additions & 3 deletions spec/lib/modules/embeddings/vector_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -94,13 +94,35 @@ def stub_vector_mapping(text, expected_embedding)
vdef.lookup_custom_param("model_name"),
text,
expected_embedding,
extra_args: {
dimensions: vdef.dimensions,
},
)
end

it_behaves_like "generates and store embeddings using a vector definition"

context "when working with models that support shortening embeddings" do
it "passes the dimensions param" do
shorter_dimensions = 10
vdef.update!(
dimensions: shorter_dimensions,
provider_params: {
model_name: "text-embedding-3-small",
},
)
text = "This is a piece of text"
short_expected_embedding = [0.0038493] * shorter_dimensions

EmbeddingsGenerationStubs.openai_service(
vdef.lookup_custom_param("model_name"),
text,
short_expected_embedding,
extra_args: {
dimensions: shorter_dimensions,
},
)

expect(described_class.new(vdef).vector_from(text)).to eq(short_expected_embedding)
end
end
end

context "with hugging_face as the provider" do
Expand Down
3 changes: 0 additions & 3 deletions spec/requests/embeddings/embeddings_controller_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -38,9 +38,6 @@ def stub_embedding(query)
vector_def.lookup_custom_param("model_name"),
query,
embedding,
extra_args: {
dimensions: vector_def.dimensions,
},
)
end

Expand Down

0 comments on commit a9d5210

Please sign in to comment.