Spaces:
Sleeping
Sleeping
File size: 11,756 Bytes
1232fcb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 |
{
"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
}
|