search

BerriAI / litellm

Call all LLM APIs using the OpenAI format. Use Bedrock, Azure, OpenAI, Cohere, Anthropic, Ollama, Sagemaker, HuggingFace, Replicate (100+ LLMs)
https://docs.litellm.ai/docs/
Other
10.05k stars 1.12k forks source link

Support for custom models imported in [Bedrock](https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-import-model.html). #4239

Closed krrishdholakia closed 1 week ago

krrishdholakia commented 1 week ago 
          Support for custom models imported in [Bedrock](https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-import-model.html).

Use case: we have a fine-tuned model deployed in Bedrock. The tuned model is based on OpenOrca, so the start and end tokens are different than instruct version.

If the provider is mistral, the template uses the instruct template rather than OpenOrca's.

  response = client.invoke_model(
      body={"prompt": "<s>[INST] hello, tell me a joke [/INST]\n", "max_tokens": 1024, "temperature": 0},
      modelId=<model_id>,
      accept=accept,
      contentType=contentType
  )

Tokens<|im_start|> and <|im_end|> should be used instead.

Tried using a custom provider as a workaround. However, the body is empty and the request fails:

  response = client.invoke_model(
      body={},
      modelId=<model_id>,
      accept=accept,
      contentType=contentType
  )

The only thing we need is that prompt template configuration is respected, as it is done with amazon or anthropic providers.

      model_id: "model_arn"
      roles: {"system":{"pre_message":"<|im_start|>system\n", "post_message":"<|im_end|>"}, "assistant":{"pre_message":"<|im_start|>assistant\n","post_message":"<|im_end|>"}, "user":{"pre_message":"<|im_start|>user\n","post_message":"<|im_end|>"}}
      bos_token: "<s>"
      eos_token: "<|im_end|>"

https://github.com/BerriAI/litellm/blob/3a35a58859a145a4a568548316a1930340e7440a/litellm/llms/bedrock.py#L743-L746

Originally posted by @andresd95 in https://github.com/BerriAI/litellm/issues/361#issuecomment-2172843397