Package: databricks-genai-inference
Databricks Generative AI Inference SDK (Beta)The Databricks Generative AI Inference Python library provides a user-friendly python interface to use the DatabricksFoundation Model API.It includes a pre-defined set of API classesEmbedding,Completion,ChatCompletionwith convenient functions to make API request, and to parse contents from raw json response.We also offer a high levelChatSessionobject for easy management of multi-round chat completions, which is especially useful for your next chatbot development.You can find more usage details in ourSDK onboarding doc.[!IMPORTANT]We're preparing to release version 1.0 of the Databricks GenerativeAI Inference Python library.Installationpipinstalldatabricks-genai-inferenceUsageEmbeddingfromdatabricks_genai_inferenceimportEmbeddingText embeddingresponse=Embedding.create(model="bge-large-en",input="3D ActionSLAM: wearable person tracking in multi-floor environments")print(f'embeddings:{response.embeddings[0]}')[!TIP]You may want to reuse http connection to improve request latency for large-scale workload, code example:withrequests.Session()asclient:fori,textinenumerate(texts):response=Embedding.create(client=client,model="bge-large-en",input=text)Text embedding (async)asyncwithhttpx.AsyncClient()asclient:response=awaitEmbedding.acreate(client=client,model="bge-large-en",input="3D ActionSLAM: wearable person tracking in multi-floor environments")print(f'embeddings:{response.embeddings[0]}')Text embedding with instructionresponse=Embedding.create(model="bge-large-en",instruction="Represent this sentence for searching relevant passages:",input="3D ActionSLAM: wearable person tracking in multi-floor environments")print(f'embeddings:{response.embeddings[0]}')Text embedding (batching)[!IMPORTANT]Support max batch size of 150response=Embedding.create(model="bge-large-en",input=["3D ActionSLAM: wearable person tracking in multi-floor environments","3D ActionSLAM: wearable person tracking in multi-floor environments"])print(f'response.embeddings[0]:{response.embeddings[0]}\n')print(f'response.embeddings[1]:{response.embeddings[1]}')Text embedding with instruction (batching)[!IMPORTANT]Support one instruction per batch
Batch sizeresponse=Embedding.create(model="bge-large-en",instruction="Represent this sentence for searching relevant passages:",input=["3D ActionSLAM: wearable person tracking in multi-floor environments","3D ActionSLAM: wearable person tracking in multi-floor environments"])print(f'response.embeddings[0]:{response.embeddings[0]}\n')print(f'response.embeddings[1]:{response.embeddings[1]}')Text completionfromdatabricks_genai_inferenceimportCompletionText completionresponse=Completion.create(model="mpt-7b-instruct",prompt="Represent the Science title:")print(f'response.text:{response.text:}')Text completion (async)asyncwithhttpx.AsyncClient()asclient:response=awaitCompletion.acreate(client=client,model="mpt-7b-instruct",prompt="Represent the Science title:")print(f'response.text:{response.text:}')Text completion (streaming)[!IMPORTANT]Only support batch size = 1 in streaming moderesponse=Completion.create(model="mpt-7b-instruct",prompt="Count from 1 to 100:",stream=True)print(f'response.text:')forchunkinresponse:print(f'{chunk.text}',end="")Text completion (streaming + async)asyncwithhttpx.AsyncClient()asclient:response=awaitCompletion.acreate(client=client,model="mpt-7b-instruct",prompt="Count from 1 to 10:",stream=True)print(f'response.text:')asyncforchunkinresponse:print(f'{chunk.text}',end="")Text completion (batching)[!IMPORTANT]Support max batch size of 16response=Completion.create(model="mpt-7b-instruct",prompt=["Represent the Science title:","Represent the Science title:"])print(f'response.text[0]:{response.text[0]}')print(f'response.text[1]:{response.text[1]}')Chat completionfromdatabricks_genai_inferenceimportChatCompletion[!IMPORTANT]Batching is not supported forChatCompletionChat completionresponse=ChatCompletion.create(model="llama-2-70b-chat",messages=[{"role":"system","content":"You are a helpful assistant."},{"role":"user","content":"Knock knock."}])print(f'response.text:{response.message:}')Chat completion (async)asyncwithhttpx.AsyncClient()asclient:response=awaitChatCompletion.acreate(client=client,model="llama-2-70b-chat",messages=[{"role":"system","content":"You are a helpful assistant."},{"role":"user","content":"Knock knock."}],)print(f'response.text:{response.message:}')Chat completion (streaming)response=ChatCompletion.create(model="llama-2-70b-chat",messages=[{"role":"system","content":"You are a helpful assistant."},{"role":"user","content":"Count from 1 to 30, add one emoji after each number"}],stream=True)forchunkinresponse:print(f'{chunk.message}',end="")Chat completion (streaming + async)asyncwithhttpx.AsyncClient()asclient:response=awaitChatCompletion.acreate(client=client,model="llama-2-70b-chat",messages=[{"role":"system","content":"You are a helpful assistant."},{"role":"user","content":"Count from 1 to 30, add one emoji after each number"}],stream=True,)asyncforchunkinresponse:print(f'{chunk.message}',end="")Chat sessionfromdatabricks_genai_inferenceimportChatSession[!IMPORTANT]Streaming mode is not supported forChatSessionchat=ChatSession(model="llama-2-70b-chat")chat.reply("Kock, kock!")print(f'chat.last:{chat.last}')chat.reply("Take a guess!")print(f'chat.last:{chat.last}')print(f'chat.history:{chat.history}')print(f'chat.count:{chat.count}')