multimodal/Titan/embeddings/v2/Titan-V2-Embeddings.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Benchmark Titan Text Embeddings V2: A new, state-of-the-art Embeddings model on Amazon Bedrock using MTEB data set\n",
    "\n",
    "Embeddings are integral to various natural language processing applications, with their quality crucial for optimal performance. They are commonly used in knowledge bases to represent textual data as dense vectors enabling efficient similarity search and retrieval. In Retrieval Augmented Generation (RAG), embeddings are used to retrieve relevant passages from a corpus to provide context for language models to generate informed, knowledge-grounded responses. Embeddings also play a key role in personalization and recommendation systems by representing user preferences, item characteristics, and historical interactions as vectors, allowing calculation of similarities for personalized recommendations based on user behavior and item embeddings. As new embedding models are released with incremental quality improvements, organizations must weigh the potential benefits against the associated costs of upgrading, considering factors like computational resources, data preprocessing, integration efforts, and projected performance gains impacting business metrics.\n",
    "\n",
    "**Some useful links for MTEB:**\n",
    "1. https://github.com/embeddings-benchmark/mteb\n",
    "\n",
    "2. https://huggingface.co/blog/mteb\n",
    "\n",
    "#### How a piece of text is converted into a vector?\n",
    "Common approach is to use models which can provide contextualized embeddings for entire sentences. These models are based on deep learning architectures such as Transformers, which can capture the contextual information and relationships between words in a sentence more effectively.\n",
    "\n",
    "![Embedding Model](./images/vector_embedding.png)\n",
    "\n",
    "In addition to semantic search, you can use embeddings to augment your prompts for more accurate results through Retrieval Augmented Generation (RAG)—but in order to use them, you’ll need to store them in a database with vector capabilities.\n",
    "\n",
    "![Embedding Model](./images/vector_db.jpg)\n",
    "\n",
    "\n",
    "In September of 2023, Amazon announced the launch of Amazon Titan Text Embeddings V1, a multilingual text embeddings model that converts text inputs like single words, phrases, or large documents into high-dimensional numerical vector representations.  Since then 1000s of our customers used the first version of the model that supported over 25 languages, with an input up to 8,192 tokens, and outputs vector of 1,536 dimensions. Today we take that to next level by introducing a flexible output embedding model. The  Amazon Titan Text Embeddings V2 model supports over 100 languages and allows for a variable dimenion output thereby saving cost in terms of reduced size to store the embeddings. The output dimensions are 256, 512 and 1024. This model is designed to perform well on multi-lingual data and use cases."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Few installs\n",
    "\n",
    "You will need to install a few libraries like MTEB for benchmark. To simply use the model you would need only boto3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install -U boto3  huggingface_hub mteb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Boto3 client to connect to the model\n",
    "\n",
    "This is optional and you can use a simple client like **boto3.client('bedrock')** to invoke. This will assume your default profile to access Bedrock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from typing import Optional\n",
    "\n",
    "# External Dependencies:\n",
    "import boto3\n",
    "from botocore.config import Config\n",
    "\n",
    "\n",
    "def get_bedrock_client(assumed_role: Optional[str] = None, region: Optional[str] = 'us-east-1',runtime: Optional[bool] = True,external_id=None, ep_url=None):\n",
    "    \"\"\"Create a boto3 client for Amazon Bedrock, with optional configuration overrides \n",
    "    \"\"\"\n",
    "    target_region = region\n",
    "\n",
    "    print(f\"Create new client\\n  Using region: {target_region}:external_id={external_id}: \")\n",
    "    session_kwargs = {\"region_name\": target_region}\n",
    "    client_kwargs = {**session_kwargs}\n",
    "\n",
    "    profile_name = os.environ.get(\"AWS_PROFILE\")\n",
    "    if profile_name:\n",
    "        print(f\"  Using profile: {profile_name}\")\n",
    "        session_kwargs[\"profile_name\"] = profile_name\n",
    "\n",
    "    retry_config = Config(\n",
    "        region_name=target_region,\n",
    "        retries={\n",
    "            \"max_attempts\": 10,\n",
    "            \"mode\": \"standard\",\n",
    "        },\n",
    "    )\n",
    "    session = boto3.Session(**session_kwargs)\n",
    "\n",
    "    if assumed_role:\n",
    "        print(f\"  Using role: {assumed_role}\", end='')\n",
    "        sts = session.client(\"sts\")\n",
    "        if external_id:\n",
    "            response = sts.assume_role(\n",
    "                RoleArn=str(assumed_role),\n",
    "                RoleSessionName=\"langchain-llm-1\",\n",
    "                ExternalId=external_id\n",
    "            )\n",
    "        else:\n",
    "            response = sts.assume_role(\n",
    "                RoleArn=str(assumed_role),\n",
    "                RoleSessionName=\"langchain-llm-1\",\n",
    "            )\n",
    "        print(f\"Using role: {assumed_role} ... sts::successful!\")\n",
    "        client_kwargs[\"aws_access_key_id\"] = response[\"Credentials\"][\"AccessKeyId\"]\n",
    "        client_kwargs[\"aws_secret_access_key\"] = response[\"Credentials\"][\"SecretAccessKey\"]\n",
    "        client_kwargs[\"aws_session_token\"] = response[\"Credentials\"][\"SessionToken\"]\n",
    "\n",
    "    if runtime:\n",
    "        service_name='bedrock-runtime'\n",
    "    else:\n",
    "        service_name='bedrock'\n",
    "\n",
    "    if ep_url:\n",
    "        bedrock_client = session.client(service_name=service_name,config=retry_config,endpoint_url = ep_url, **client_kwargs )\n",
    "    else:\n",
    "        bedrock_client = session.client(service_name=service_name,config=retry_config, **client_kwargs )\n",
    "\n",
    "    print(\"boto3 Bedrock client successfully created!\")\n",
    "    print(bedrock_client._endpoint)\n",
    "    return bedrock_client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Helper Class to connect and run the embeddings\n",
    "\n",
    "This is not a production use code but more of a reference "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import boto3\n",
    "\n",
    "class TitanEmbeddings(object):\n",
    "    accept = \"application/json\"\n",
    "    content_type = \"application/json\"\n",
    "    \n",
    "    def __init__(self, model_id=\"amazon.titan-embed-text-v2:0\", boto3_client=None, region_name='us-east-1'):\n",
    "        \n",
    "        if boto3_client:\n",
    "            self.bedrock_boto3 = boto3_client\n",
    "        else:\n",
    "            # self.bedrock_boto3 = boto3.client(service_name='bedrock-runtime')\n",
    "            self.bedrock_boto3 = boto3.client(\n",
    "                service_name='bedrock-runtime', \n",
    "                region_name=region_name, \n",
    "            )\n",
    "        self.model_id = model_id\n",
    "\n",
    "    def __call__(self, text, dimensions, normalize=True):\n",
    "        \"\"\"\n",
    "        Returns Titan Embeddings\n",
    "\n",
    "        Args:\n",
    "            text (str): text to embed\n",
    "            dimensions (int): Number of output dimensions.\n",
    "            normalize (bool): Whether to return the normalized embedding or not.\n",
    "\n",
    "        Return:\n",
    "            List[float]: Embedding\n",
    "            \n",
    "        \"\"\"\n",
    "\n",
    "        body = json.dumps({\n",
    "            \"inputText\": text,\n",
    "            \"dimensions\": dimensions,\n",
    "            \"normalize\": normalize\n",
    "        })\n",
    "\n",
    "        response = self.bedrock_boto3.invoke_model(\n",
    "            body=body, modelId=self.model_id, accept=self.accept, contentType=self.content_type\n",
    "        )\n",
    "\n",
    "        response_body = json.loads(response.get('body').read())\n",
    "\n",
    "        return response_body['embedding']\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test connection to bedrock\n",
    "\n",
    "use the bedrock boto client. If you are using a role to be assumed pass that in. If you have the profile set up then leverage by setting it in the os environ variable AWS_PROFILE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import sys\n",
    "\n",
    "import boto3\n",
    "\n",
    "boto3_bedrock_runtime = get_bedrock_client() #boto3.client('bedrock')\n",
    "\n",
    "bedrock_embeddings = TitanEmbeddings(model_id=\"amazon.titan-embed-text-v2:0\", boto3_client=boto3_bedrock_runtime)\n",
    "bedrock_embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate Embeddings\n",
    "\n",
    "At the time of writing you can use amazon.titan-embed-text-v2 as embedding model via the API. The input text size is 8k tokens and the output vector length can be any of 256, 512 or 1024\n",
    "\n",
    "To use a text embeddings model, use the InvokeModel API operation or the Python SDK. Use InvokeModel to retrieve the vector representation of the input text from the specified model.\n",
    "Input\n",
    "\n",
    "```\n",
    "\n",
    "{\n",
    "    \"inputText\": text,\n",
    "    \"dimensions\": dimensions, # range from 256 , 512, 1024\n",
    "    \"normalize\": normalize\n",
    "}\n",
    "\n",
    "Output\n",
    "\n",
    "{\n",
    "    \"embedding\": []\n",
    "}\n",
    "```\n",
    "\n",
    "#### Normalization of a vector \n",
    "\n",
    "Normalization is the process of scaling it to have a unit length or magnitude of 1. It is useful to ensure that all vectors have the same scale and contribute equally during vector operations, preventing some vectors from dominating others due to their larger magnitudes.\n",
    "\n",
    "#### When should you Normalize:\n",
    "Use this as default for most of the use cases like Retrieval, RAG and others\n",
    "\n",
    "#### When you should not Normalize: \n",
    "Normnally normalization wil work for all use cases, but you can experiment for certain use cases like Classification or Entity extraction\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "prompt_data = \"Amazon Bedrock supports foundation models from industry-leading providers such as \\\n",
    "AI21 Labs, Anthropic, Stability AI, and Amazon. Choose the model that is best suited to achieving \\\n",
    "your unique goals.\"\n",
    "\n",
    "\n",
    "modelId = \"amazon.titan-embed-text-v2:0\"  # \n",
    "accept = \"application/json\"\n",
    "contentType = \"application/json\"\n",
    "\n",
    "\n",
    "\n",
    "sample_model_input={\n",
    "    \"inputText\": prompt_data,\n",
    "    \"dimensions\": 256,\n",
    "    \"normalize\": True\n",
    "}\n",
    "\n",
    "body = json.dumps(sample_model_input)\n",
    "\n",
    "response = boto3_bedrock_runtime.invoke_model(body=body, modelId=modelId, accept=accept, contentType=contentType)\n",
    "\n",
    "response_body = json.loads(response.get('body').read())\n",
    "\n",
    "embedding = response_body.get(\"embedding\")\n",
    "print(f\"The embedding vector has {len(embedding)} values\\n{embedding[0:3]+['...']+embedding[-3:]}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Use the Helper class to show the same results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "embedding= bedrock_embeddings(text=prompt_data, dimensions=256, normalize=True)\n",
    "print(f\"The embedding vector has {len(embedding)} values\\n{embedding[0:3]+['...']+embedding[-3:]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optional MTEB testing\n",
    "\n",
    "To create your own MTEB testing you can use this code below as a starting point and a sample. Please note this is not a production code. You should refer to the github location for MTEB at https://github.com/embeddings-benchmark/mteb/tree/main for more details"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import json\n",
    "import os\n",
    "import sys\n",
    "import time\n",
    "import torch\n",
    "\n",
    "import boto3\n",
    "\n",
    "class TitanV2Model():\n",
    "    def __init__(self) -> None:\n",
    "\n",
    "        self.br_embeddings = None     \n",
    "        self._init_connection()\n",
    "        \n",
    "        print(f\"TitanV2Model:__init__::ready:to:Invoke:::successful::\") \n",
    "    \n",
    "    def _init_connection(self, dim=256):\n",
    "        boto3_bedrock_runtime = get_bedrock_client() #boto3.client('bedrock')\n",
    "\n",
    "        self.br_embeddings = TitanEmbeddings(model_id=\"amazon.titan-embed-text-v2:0\", boto3_client=boto3_bedrock_runtime)\n",
    "        self.dim = dim\n",
    "        \n",
    "    def process_dict_text(self, single_text_dict):\n",
    "        \"\"\" **IMPORTANT** CHANGE this Code to be tuned to your data set and use this -- DO NOT USE THIS AS IS. refer to  https://github.com/embeddings-benchmark/mteb/blob/main/mteb/abstasks/AbsTaskRetrieval.py as an example. Please goover this Git hub in detail\"\"\"\n",
    "        single_text = [str(key).strip() + \" \" + str(val).strip() if 'title' in key.lower() else str(val).strip() for key, val in single_text_dict.items()]\n",
    "        return \" \".join(single_text)[:30000]\n",
    "        \n",
    "    def reorg_text(self, single_text):\n",
    "        \n",
    "        \"\"\" **IMPORTANT** CHANGE this Code to be tuned to your data set and use this -- DO NOT USE THIS AS IS. refer to  https://github.com/embeddings-benchmark/mteb/blob/main/mteb/abstasks/AbsTaskRetrieval.py as an example. Please goover this Git hub in detail\"\"\"\n",
    "        if isinstance(single_text, dict):\n",
    "            single_text = self.process_dict_text(single_text)\n",
    "        single_text = \"0\" if not single_text else single_text \n",
    "        # check for json -- \n",
    "        try:\n",
    "            single_text_dict = json.loads(single_text)\n",
    "            single_text = self.process_dict_text(single_text_dict)\n",
    "        except:\n",
    "            pass\n",
    "        return single_text\n",
    "    \n",
    "    def invoke_model(self, text_list: list[str]):\n",
    "        \"\"\" **IMPORTANT** CHANGE this Code to be tuned to your data set and use this -- DO NOT USE THIS AS IS. refer to  https://github.com/embeddings-benchmark/mteb/blob/main/mteb/abstasks/AbsTaskRetrieval.py as an example. Please goover this Git hub in detail\"\"\"\n",
    "        list_embeddings = []\n",
    "        \n",
    "        for single_text in text_list:\n",
    "            single_text = self.reorg_text(single_text)\n",
    "            single_embed = bedrock_embeddings(text=single_text, dimensions=self.dim, normalize=True)\n",
    "            list_embeddings.append(single_embed)\n",
    "\n",
    "        return list_embeddings\n",
    "\n",
    "    def reshape_titan_embeddings(self, query_embeddings: np.ndarray, **kwargs) -> list[np.ndarray]:\n",
    "        # - use this to re shape your embeddings as needed\n",
    "        return query_embeddings # \n",
    "        \n",
    "        \n",
    "    def encode(self, queries: list[str], **kwargs) -> list[np.ndarray] | list[torch.Tensor] : # - | list[torch.Tensor] \n",
    "        \"\"\"\n",
    "        Returns a list of embeddings for the given sentences.\n",
    "        Args:\n",
    "            queries: List of sentences to encode\n",
    "\n",
    "        Returns:\n",
    "            List of embeddings for the given sentences\n",
    "        \"\"\"\n",
    "        \n",
    "        embedding_list = self.invoke_model(queries)\n",
    "        return self.reshape_titan_embeddings(np.array(embedding_list))\n",
    "\n",
    "    \n",
    "    def encode_queries(self, queries: list[str], **kwargs) -> list[np.ndarray] | list[torch.Tensor] : # - | list[torch.Tensor] \n",
    "        \"\"\"\n",
    "        Returns a list of embeddings for the given sentences.\n",
    "        Args:\n",
    "            queries: List of sentences to encode\n",
    "\n",
    "        Returns:\n",
    "            List of embeddings for the given sentences\n",
    "        \"\"\"\n",
    "        \n",
    "        embedding_list = self.invoke_model(queries)\n",
    "        return self.reshape_titan_embeddings(np.array(embedding_list))\n",
    "\n",
    "\n",
    "    def encode_corpus(self, corpus: list[str] | list[dict[str, str]], **kwargs) -> list[np.ndarray] | list[torch.Tensor] : #- | list[torch.Tensor]:\n",
    "        \"\"\"\n",
    "        Returns a list of embeddings for the given sentences.\n",
    "        Args:\n",
    "            corpus: List of sentences to encode\n",
    "                or list of dictionaries with keys \"title\" and \"text\"\n",
    "\n",
    "        Returns:\n",
    "            List of embeddings for the given sentences\n",
    "        \"\"\"\n",
    "        \n",
    "        embedding_list = self.invoke_model(corpus)\n",
    "        return self.reshape_titan_embeddings(np.array(embedding_list))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "final_output = TitanV2Model().encode_queries([\"this is a test\"])\n",
    "\n",
    "print(f\"Embneddings Generated ::\")\n",
    "print(f\"shape:of:embeddings -- > length of embeddings={len(final_output)}::\")\n",
    "print(f\"shape:of:embeddings -- > {len(final_output[0])}::\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sample of one of the classfication tasks which can be run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from mteb import MTEB\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from mteb.tasks.Retrieval.eng.CQADupstackEnglishRetrieval import CQADupstackEnglishRetrieval \n",
    "\n",
    "model = TitanV2Model()\n",
    "evaluation = MTEB(tasks=[CQADupstackEnglishRetrieval(langs=[\"en\"]),])\n",
    "\n",
    "evaluation.run(model, eval_splits=[\"test\"], trust_remote_code=True, overwrite_results = True) #- output_folder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Please refer to this link [for Amazon Titan models](https://docs.aws.amazon.com/bedrock/latest/userguide/titan-models.html) "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p310",
   "language": "python",
   "name": "conda_pytorch_p310"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}