ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

(feat) add ollama/llava

Browse Source
This commit is contained in:
ishaan-jaff 2023-12-16 10:34:20 +05:30
parent add153d110
commit 287633887e
2 changed files with 27 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
litellm/llms/prompt_templates/factory.py
View File

@ -73,6 +73,25 @@ def ollama_pt(model, messages): # https://github.com/jmorganca/ollama/blob/af4cf
final_prompt_value="### Response:",
messages=messages
)
elif "llava" in model:
prompt = ""
images = []
for message in messages:
if isinstance(message["content"], str):
prompt += message["content"]
elif isinstance(message["content"], list):
# see https://docs.litellm.ai/docs/providers/openai#openai-vision-models
for element in message["content"]:
if isinstance(element, dict):
if element["type"] == "text":
prompt += element["text"]
elif element["type"] == "image_url":
image_url = element["image_url"]["url"]
images.append(image_url)
return {
"prompt": prompt,
"images": images
}
else:
prompt = "".join(m["content"] if isinstance(m['content'], str) is str else "".join(m['content']) for m in messages)
return prompt

9
litellm/main.py
View File

@ -1308,7 +1308,14 @@ def completion(
)
else:
prompt = prompt_factory(model=model, messages=messages, custom_llm_provider=custom_llm_provider)
if isinstance(prompt, dict):
# for multimode models - ollama/llava prompt_factory returns a dict {
# "prompt": prompt,
# "images": images
# }
prompt, images = prompt["prompt"], prompt["images"]
optional_params["images"] = images
## LOGGING
generator = ollama.get_ollama_response_stream(api_base, model, prompt, optional_params, logging_obj=logging, acompletion=acompletion, model_response=model_response, encoding=encoding)
if acompletion is True: