pref: change ollama embedded api request (#6876)

This commit is contained in:
灰灰 2024-08-02 12:04:47 +08:00 committed by GitHub
parent f8617db012
commit 56af1a0adf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -59,7 +59,7 @@ class OllamaEmbeddingModel(TextEmbeddingModel):
if not endpoint_url.endswith('/'): if not endpoint_url.endswith('/'):
endpoint_url += '/' endpoint_url += '/'
endpoint_url = urljoin(endpoint_url, 'api/embeddings') endpoint_url = urljoin(endpoint_url, 'api/embed')
# get model properties # get model properties
context_size = self._get_context_size(model, credentials) context_size = self._get_context_size(model, credentials)
@ -78,12 +78,9 @@ class OllamaEmbeddingModel(TextEmbeddingModel):
else: else:
inputs.append(text) inputs.append(text)
batched_embeddings = []
for text in inputs:
# Prepare the payload for the request # Prepare the payload for the request
payload = { payload = {
'prompt': text, 'input': inputs,
'model': model, 'model': model,
} }
@ -99,11 +96,10 @@ class OllamaEmbeddingModel(TextEmbeddingModel):
response_data = response.json() response_data = response.json()
# Extract embeddings and used tokens from the response # Extract embeddings and used tokens from the response
embeddings = response_data['embedding'] embeddings = response_data['embeddings']
embedding_used_tokens = self.get_num_tokens(model, credentials, [text]) embedding_used_tokens = self.get_num_tokens(model, credentials, inputs)
used_tokens += embedding_used_tokens used_tokens += embedding_used_tokens
batched_embeddings.append(embeddings)
# calc usage # calc usage
usage = self._calc_response_usage( usage = self._calc_response_usage(
@ -113,7 +109,7 @@ class OllamaEmbeddingModel(TextEmbeddingModel):
) )
return TextEmbeddingResult( return TextEmbeddingResult(
embeddings=batched_embeddings, embeddings=embeddings,
usage=usage, usage=usage,
model=model model=model
) )