{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AWS Setup\n",
    "\n",
    "1. Get your AWS access key and secret key from AWS console\n",
    "   1. Log in to the AWS Management Console (it is OK if you see some Access denied messages on the front page. The account is limited on purpose)\n",
    "   2. Click on your name at the top-right coder and then \"Security Credentials\"\n",
    "   3. Click on \"Access keys\" and create a new access key. Create an access key for the Command Line Interface (CLI).\n",
    "   4. Download and save your access key and secret key\n",
    "2. Install the AWS CLI tool (optional but recommended): https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html\n",
    "3. Configure the AWS CLI tool to use your access key and secret key:\n",
    "    ```bash\n",
    "    aws configure --profile sigir-participant\n",
    "    # Use the following settings:\n",
    "    # AWS Access Key ID\n",
    "    # AWS Secret Access Key\n",
    "    # Default region name: us-east-1\n",
    "    ```\n",
    "4. Test your setup by running the following command:\n",
    "    ```bash\n",
    "    # Should display your AWS account ID\n",
    "    aws sts get-caller-identity --profile sigir-participant\n",
    "\n",
    "    # Make sure you are able to access the configuraiton service\n",
    "    aws ssm get-parameter --name /pinecone/ro_token --profile sigir-participant\n",
    "\n",
    "    ```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Python Setup\n",
    "1. Install an up-to-date python version (we use 3.12, but any recent version should work)\n",
    "2. Install the required python packages (see below). We recommend using a virtual environment, e.g. venv or conda or poetry or uv\n",
    "3. Now you are ready to try the code samples below\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install dependencies\n",
    "# The following package versions were tested using Python 3.12 and proved to work but it is possible to use\n",
    "# different versions, consider them simply as examples.\n",
    "\n",
    "!pip install torch==2.5.1 transformers==4.45.2 boto3==1.35.88 pinecone==5.4.2 opensearch-py==2.8.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# AWS utilities\n",
    "import boto3\n",
    "\n",
    "AWS_PROFILE_NAME = \"sigir-participant\"\n",
    "AWS_REGION_NAME = \"us-east-1\"\n",
    "\n",
    "def get_ssm_value(key: str, profile: str = AWS_PROFILE_NAME, region: str = AWS_REGION_NAME) -> str:\n",
    "    \"\"\"Get a cleartext value from AWS SSM.\"\"\"\n",
    "    session = boto3.Session(profile_name=profile, region_name=region)\n",
    "    ssm = session.client(\"ssm\")\n",
    "    return ssm.get_parameter(Name=key)[\"Parameter\"][\"Value\"]\n",
    "\n",
    "def get_ssm_secret(key: str, profile: str = AWS_PROFILE_NAME, region: str = AWS_REGION_NAME):\n",
    "    session = boto3.Session(profile_name=profile, region_name=region)\n",
    "    ssm = session.client(\"ssm\")\n",
    "    return ssm.get_parameter(Name=key, WithDecryption=True)[\"Parameter\"][\"Value\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pinecone sample\n",
    "\n",
    "from typing import List, Literal, Tuple\n",
    "from multiprocessing.pool import ThreadPool\n",
    "import boto3\n",
    "from pinecone import Pinecone\n",
    "import torch\n",
    "from functools import cache\n",
    "from transformers import AutoModel, AutoTokenizer\n",
    "\n",
    "PINECONE_INDEX_NAME = \"fineweb10bt-512-0w-e5-base-v2\"\n",
    "PINECONE_NAMESPACE=\"default\"\n",
    "\n",
    "@cache\n",
    "def has_mps():\n",
    "    return torch.backends.mps.is_available()\n",
    "\n",
    "@cache\n",
    "def has_cuda():\n",
    "    return torch.cuda.is_available()\n",
    "\n",
    "@cache\n",
    "def get_tokenizer(model_name: str = \"intfloat/e5-base-v2\"):\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "    return tokenizer\n",
    "\n",
    "@cache\n",
    "def get_model(model_name: str = \"intfloat/e5-base-v2\"):\n",
    "    model = AutoModel.from_pretrained(model_name, trust_remote_code=True)\n",
    "    if has_mps():\n",
    "        model = model.to(\"mps\")\n",
    "    elif has_cuda():\n",
    "        model = model.to(\"cuda\")\n",
    "    else:\n",
    "        model = model.to(\"cpu\")\n",
    "    return model\n",
    "\n",
    "def average_pool(last_hidden_states: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:\n",
    "    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)\n",
    "    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]\n",
    "\n",
    "def embed_query(query: str,\n",
    "                query_prefix: str = \"query: \",\n",
    "                model_name: str = \"intfloat/e5-base-v2\",\n",
    "                pooling: Literal[\"cls\", \"avg\"] = \"avg\",\n",
    "                normalize: bool =True) -> list[float]:\n",
    "    return batch_embed_queries([query], query_prefix, model_name, pooling, normalize)[0]\n",
    "\n",
    "def batch_embed_queries(queries: List[str], query_prefix: str = \"query: \", model_name: str = \"intfloat/e5-base-v2\", pooling: Literal[\"cls\", \"avg\"] = \"avg\", normalize: bool =True) -> List[List[float]]:\n",
    "    with_prefixes = [\" \".join([query_prefix, query]) for query in queries]\n",
    "    tokenizer = get_tokenizer(model_name)\n",
    "    model = get_model(model_name)\n",
    "    with torch.no_grad():\n",
    "        encoded = tokenizer(with_prefixes, padding=True, return_tensors=\"pt\", truncation=\"longest_first\")\n",
    "        encoded = encoded.to(model.device)\n",
    "        model_out = model(**encoded)\n",
    "        match pooling:\n",
    "            case \"cls\":\n",
    "                embeddings = model_out.last_hidden_state[:, 0]\n",
    "            case \"avg\":\n",
    "                embeddings = average_pool(model_out.last_hidden_state, encoded[\"attention_mask\"])\n",
    "        if normalize:\n",
    "            embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)\n",
    "    return embeddings.tolist()\n",
    "\n",
    "@cache\n",
    "def get_pinecone_index(index_name: str = PINECONE_INDEX_NAME):\n",
    "    pc = Pinecone(api_key=get_ssm_secret(\"/pinecone/ro_token\"))\n",
    "    index = pc.Index(name=index_name)\n",
    "    return index\n",
    "\n",
    "def query_pinecone(query: str, top_k: int = 10, namespace: str = PINECONE_NAMESPACE) -> dict:\n",
    "    index = get_pinecone_index()\n",
    "    results = index.query(\n",
    "        vector=embed_query(query),\n",
    "        top_k=top_k,\n",
    "        include_values=False,\n",
    "        namespace=namespace,\n",
    "        include_metadata=True\n",
    "    )\n",
    "\n",
    "    return results\n",
    "\n",
    "def batch_query_pinecone(queries: list[str], top_k: int = 10, namespace: str = PINECONE_NAMESPACE, n_parallel: int = 10) -> list[dict]:\n",
    "    \"\"\"Batch query a Pinecone index and return the results.\n",
    "\n",
    "    Internally uses a ThreadPool to parallelize the queries.\n",
    "    \"\"\"\n",
    "    index = get_pinecone_index()\n",
    "    embeds = batch_embed_queries(queries)\n",
    "    pool = ThreadPool(n_parallel)\n",
    "    results = pool.map(lambda x: index.query(vector=x, top_k=top_k, include_values=False, namespace=namespace, include_metadata=True), embeds)\n",
    "    return results\n",
    "\n",
    "def show_pinecone_results(results):\n",
    "    for match in results[\"matches\"]:\n",
    "        print(\"chunk:\", match[\"id\"], \"score:\", match[\"score\"])\n",
    "        print(match[\"metadata\"][\"text\"])\n",
    "        print()\n",
    "\n",
    "results = query_pinecone(\"What is a second brain?\")\n",
    "show_pinecone_results(results)\n",
    "\n",
    "batch_results = batch_query_pinecone([\"What is a second brain?\", \"how does a brain work?\", \"Where is Paris?\"], top_k=2)\n",
    "for results in batch_results:\n",
    "    show_pinecone_results(results)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# OpenSearch sample\n",
    "from functools import cache\n",
    "from opensearchpy import OpenSearch, AWSV4SignerAuth, RequestsHttpConnection\n",
    "\n",
    "OPENSEARCH_INDEX_NAME = \"fineweb10bt-512-0w-e5-base-v2\"\n",
    "\n",
    "@cache\n",
    "def get_client(profile: str = AWS_PROFILE_NAME, region: str = AWS_REGION_NAME):\n",
    "    credentials = boto3.Session(profile_name=profile).get_credentials()\n",
    "    auth = AWSV4SignerAuth(credentials, region=region)\n",
    "    host_name = get_ssm_value(\"/opensearch/endpoint\", profile=profile, region=region)\n",
    "    aos_client = OpenSearch(\n",
    "        hosts=[{\"host\": host_name, \"port\": 443}],\n",
    "        http_auth=auth,\n",
    "        use_ssl=True,\n",
    "        verify_certs=True,\n",
    "        connection_class=RequestsHttpConnection,\n",
    "    )\n",
    "    return aos_client\n",
    "\n",
    "def query_opensearch(query: str, top_k: int = 10) -> dict:\n",
    "    \"\"\"Query an OpenSearch index and return the results.\"\"\"\n",
    "    client = get_client()\n",
    "    results = client.search(index=OPENSEARCH_INDEX_NAME, body={\"query\": {\"match\": {\"text\": query}}, \"size\": top_k})\n",
    "    return results\n",
    "\n",
    "def batch_query_opensearch(queries: list[str], top_k: int = 10, n_parallel: int = 10) -> list[dict]:\n",
    "    \"\"\"Sends a list of queries to OpenSearch and returns the results.",
    " Configuration of Connection Timeout might be needed for serving large batches of queries",
    "\"\"\"\n",
    "    client = get_client()\n",
    "    request = []\n",
    "    for query in queries:\n",
    "        req_head = {\"index\": OPENSEARCH_INDEX_NAME}\n",
    "        req_body = {\n",
    "            \"query\": {\n",
    "                \"multi_match\": {\n",
    "                    \"query\": query,\n",
    "                    \"fields\": [\"text\"],\n",
    "                }\n",
    "            },\n",
    "            \"size\": top_k,\n",
    "        }\n",
    "        request.extend([req_head, req_body])\n",
    "\n",
    "    return client.msearch(body=request)\n",
    "\n",
    "\n",
    "\n",
    "def show_opensearch_results(results: dict):\n",
    "    for match in results[\"hits\"][\"hits\"]:\n",
    "        print(\"chunk:\", match[\"_id\"], \"score:\", match[\"_score\"])\n",
    "        print(match[\"_source\"][\"text\"])\n",
    "        print()\n",
    "\n",
    "results = query_opensearch(\"What is a second brain?\")\n",
    "show_opensearch_results(results)",
    "\n",
    "batch_results = batch_query_opensearch([\"What is a second brain?\", \"how does a brain work?\", \"Where is Paris?\"], top_k=1)\n",
    "\n",
    "for results in batch_results['responses']:\n",
    "    show_opensearch_results(results)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}