phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

Merge pull request #2720 from onukura/ollama-batch-embedding

Browse source 
Batch embedding for Ollama
This commit is contained in:
Krish Dholakia 2024-03-28 14:58:55 -07:00 • committed by GitHub
commit 28905c85b6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 52 additions and 62 deletions
Download patch file Download diff file Expand all files Collapse all files

34
litellm/llms/ollama.py
View file

@ -344,9 +344,9 @@ async def ollama_acompletion(url, data, model_response, encoding, logging_obj):
async def ollama_aembeddings(
api_base="http://localhost:11434",
model="llama2",
prompt="Why is the sky blue?",
api_base: str,
model: str,
prompts: list[str],
optional_params=None,
logging_obj=None,
model_response=None,
@ -365,6 +365,11 @@ async def ollama_aembeddings(
): # completion(top_k=3) > cohere_config(top_k=3) <- allows for dynamic variables to be passed in
optional_params[k] = v
total_input_tokens = 0
output_data = []
timeout = aiohttp.ClientTimeout(total=litellm.request_timeout) # 10 minutes
async with aiohttp.ClientSession(timeout=timeout) as session:
for idx, prompt in enumerate(prompts):
data = {
"model": model,
"prompt": prompt,
@ -375,10 +380,8 @@ async def ollama_aembeddings(
api_key=None,
additional_args={"api_base": url, "complete_input_dict": data, "headers": {}},
)
timeout = aiohttp.ClientTimeout(total=litellm.request_timeout) # 10 minutes
async with aiohttp.ClientSession(timeout=timeout) as session:
response = await session.post(url, json=data)
response = await session.post(url, json=data)
if response.status != 200:
text = await response.text()
raise OllamaError(status_code=response.status, message=text)
@ -395,22 +398,19 @@ async def ollama_aembeddings(
)
response_json = await response.json()
embeddings = response_json["embedding"]
embeddings = [embeddings] # Ollama currently does not support batch embedding
## RESPONSE OBJECT
output_data = []
for idx, embedding in enumerate(embeddings):
embeddings: list[float] = response_json["embedding"]
output_data.append(
{"object": "embedding", "index": idx, "embedding": embedding}
{"object": "embedding", "index": idx, "embedding": embeddings}
)
input_tokens = len(encoding.encode(prompt))
total_input_tokens += input_tokens
model_response["object"] = "list"
model_response["data"] = output_data
model_response["model"] = model
input_tokens = len(encoding.encode(prompt))
model_response["usage"] = {
"prompt_tokens": input_tokens,
"total_tokens": input_tokens,
"prompt_tokens": total_input_tokens,
"total_tokens": total_input_tokens,
}
return model_response

20
litellm/main.py
View file

@ -2796,29 +2796,19 @@ def embedding(
or get_secret("OLLAMA_API_BASE")
or "http://localhost:11434"
)
ollama_input = None
if isinstance(input, list) and len(input) > 1:
raise litellm.BadRequestError(
message=f"Ollama Embeddings don't support batch embeddings",
model=model, # type: ignore
llm_provider="ollama", # type: ignore
)
if isinstance(input, list) and len(input) == 1:
ollama_input = "".join(input[0])
elif isinstance(input, str):
ollama_input = input
else:
if isinstance(input ,str):
input = [input]
if not all(isinstance(item, str) for item in input):
raise litellm.BadRequestError(
message=f"Invalid input for ollama embeddings. input={input}",
model=model, # type: ignore
llm_provider="ollama", # type: ignore
)
if aembedding == True:
if aembedding:
response = ollama.ollama_aembeddings(
api_base=api_base,
model=model,
prompt=ollama_input,
prompts=input,
encoding=encoding,
logging_obj=logging,
optional_params=optional_params,