{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Databricks Embeddings\n",
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index\n",
    "%pip install llama-index-embeddings-databricks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from llama_index.core import Settings\n",
    "from llama_index.embeddings.databricks import DatabricksEmbedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the DatabricksEmbedding class with the required model, API key and serving endpoint\n",
    "os.environ[\"DATABRICKS_TOKEN\"] = \"<MY TOKEN>\"\n",
    "os.environ[\"DATABRICKS_SERVING_ENDPOINT\"] = \"<MY ENDPOINT>\"\n",
    "embed_model = DatabricksEmbedding(model=\"databricks-bge-large-en\")\n",
    "Settings.embed_model = embed_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Embed some text\n",
    "embeddings = embed_model.get_text_embedding(\n",
    "    \"The DatabricksEmbedding integration works great.\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}