<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Low-latency item-to-item recommendation system - Creating ANN index
Overview
This notebook is a part of the series that describes the process of implementing a Low-latency item-to-item recommendation system.
The notebook demonstrates how to create and deploy an ANN index using item embeddings created in the preceding notebooks. In this notebook you go through the following steps.
Exporting embeddings from BigQuery into the JSONL formated file.
Creating an ANN Index using the exported embeddings.
Creating and ANN Endpoint. 
Deploying the ANN Index to the ANN Endpoint.
Testing the deployed ANN Index.
This notebook was designed to run on AI Platform Notebooks. Before running the notebook make sure that you have completed the setup steps as described in the README file.
Setting up the notebook's environment
Import notebook dependencies
Step1:  In the experimental release, the Online Querying API of the ANN service is exposed throught the GRPC interface. The ann_grpc folder contains the grpc client stub to interface to the API.
Step2:  Configure GCP environment
Set the following constants to the values reflecting your environment
Step4:  Exporting the embeddings
In the preceeding notebooks you trained the Matrix Factorization BQML model and exported the embeddings to the item_embeddings table. 
In this step you will extract the embeddings to a set of JSONL files in the format required by the ANN service.
Verify the number of embeddings
Step5:  Export the embeddings
You will use the BigQuery export job to export the embeddings table.
Step6:  Inspect the extracted files.
Step20:  Creating an ANN index deployment
Deploying an ANN index is a 3 step process
Step21:  Create an ANN index
List all indexes registered with the ANN service
Step22:  Configure and create a new index based on the exported embeddings
Index creation is a long running operation. Be patient.
Step23:  Verify that the index was created
Step24:  Create the index deployment
List all endpoints registered with the ANN service
Step25:  Create an index endpoint
Step26:  Verify that the endpoint was created
Step27:  Deploy the index to the endpoint
Set the deployed index ID
The ID of the deployed index must be unique within your project.
Step28:  Deploy the index
Be patient. Index deployment is a long running operation
Step29:  Querying the ANN service
You will use the gRPC interface to query the deployed index.
Retrieve the gRPC private endpoint for the ANN Match service
Step33:  Create a helper wrapper around the Match Service gRPC API.
The wrapper uses the pre-generated gRPC stub to the Online Querying gRPC interface.
Step34:  Prepare sample data
Retrieve a few embeddings from the BigQuery embedding table
Step35:  Run a single match query
The following call requests 10 closest neighbours for a single embedding.
Step36:  Run a batch match query
The following call requests 3 closest neighbours for each of the embeddings in a batch of 5.
Step37:  Clean up
WARNING
The below code will delete all ANN deployments, endpoints, and indexes in the configured project.
Delete index deployments and endpoints
Step38:  Delete indexes
<ASSISTANT_TASK:>
Python Code:
import base64
import datetime
import logging
import os
import json
import pandas as pd
import time
import sys
import grpc
import google.auth
import numpy as np
import tensorflow.io as tf_io
from google.cloud import bigquery
from typing import List, Optional, Text, Tuple
Explanation: Low-latency item-to-item recommendation system - Creating ANN index
Overview
This notebook is a part of the series that describes the process of implementing a Low-latency item-to-item recommendation system.
The notebook demonstrates how to create and deploy an ANN index using item embeddings created in the preceding notebooks. In this notebook you go through the following steps.
Exporting embeddings from BigQuery into the JSONL formated file.
Creating an ANN Index using the exported embeddings.
Creating and ANN Endpoint. 
Deploying the ANN Index to the ANN Endpoint.
Testing the deployed ANN Index.
This notebook was designed to run on AI Platform Notebooks. Before running the notebook make sure that you have completed the setup steps as described in the README file.
Setting up the notebook's environment
Import notebook dependencies
End of explanation
ANN_GRPC_ENDPOINT_STUB = 'ann_grpc'
if ANN_GRPC_ENDPOINT_STUB not in sys.path:
    sys.path.append(ANN_GRPC_ENDPOINT_STUB)
import ann_grpc.match_pb2_grpc as match_pb2_grpc
import ann_grpc.match_pb2 as match_pb2
Explanation: In the experimental release, the Online Querying API of the ANN service is exposed throught the GRPC interface. The ann_grpc folder contains the grpc client stub to interface to the API.
End of explanation
PROJECT_ID = 'jk-mlops-dev' # <-CHANGE THIS
PROJECT_NUMBER = '895222332033' # <-CHANGE THIS
BQ_DATASET_NAME = 'song_embeddings' # <- CHANGE THIS
BQ_LOCATION = 'US' # <- CHANGE THIS
DATA_LOCATION = 'gs://jk-ann-staging/embeddings' # <-CHANGE THIS
VPC_NAME = 'default' # <-CHANGE THIS
EMBEDDINGS_TABLE = 'item_embeddings'
REGION = 'us-central1'
MATCH_SERVICE_PORT = 10000
Explanation: Configure GCP environment
Set the following constants to the values reflecting your environment:
PROJECT_ID - your GCP project ID.
PROJECT_NUMBER - your GCP project number.
BQ_DATASET_NAME - the name of the BigQuery dataset that contains the item embeddings table.
BQ_LOCATION - the dataset location
DATA_LOCATION - a GCS location for the exported embeddings (JSONL) files.
VPC_NAME - a name of the GCP VPC to use for the index deployments. Use the name of the VPC prepared during the initial setup. 
REGION - a compute region. Don't change the default - us-central - while the ANN Service is in the experimental stage
End of explanation
client = bigquery.Client(project=PROJECT_ID, location=BQ_LOCATION)
query = f
    SELECT COUNT(*) embedding_count
    FROM {BQ_DATASET_NAME}.item_embeddings;
query_job = client.query(query)
query_job.to_dataframe()
Explanation: Exporting the embeddings
In the preceeding notebooks you trained the Matrix Factorization BQML model and exported the embeddings to the item_embeddings table. 
In this step you will extract the embeddings to a set of JSONL files in the format required by the ANN service.
Verify the number of embeddings
End of explanation
file_name_pattern = 'embedding-*.json'
destination_uri = f'{DATA_LOCATION}/{file_name_pattern}'
table_id = 'item_embeddings'
destination_format = 'NEWLINE_DELIMITED_JSON'
dataset_ref = bigquery.DatasetReference(PROJECT_ID, BQ_DATASET_NAME)
table_ref = dataset_ref.table(table_id)
job_config = bigquery.job.ExtractJobConfig()
job_config.destination_format = bigquery.DestinationFormat.NEWLINE_DELIMITED_JSON
extract_job = client.extract_table(
    table_ref,
    destination_uris=destination_uri,
    job_config=job_config,
    #location=BQ_LOCATION,
)  
extract_job.result()
Explanation: Export the embeddings
You will use the BigQuery export job to export the embeddings table.
End of explanation
! gsutil ls {DATA_LOCATION}
Explanation: Inspect the extracted files.
End of explanation
import datetime
import logging
import json
import time
import google.auth
class ANNClient(object):
    Base ANN Service client.
    
    def __init__(self, project_id, project_number, region):
        credentials, _ = google.auth.default()
        self.authed_session = google.auth.transport.requests.AuthorizedSession(credentials)
        self.ann_endpoint = f'{region}-aiplatform.googleapis.com'
        self.ann_parent = f'https://{self.ann_endpoint}/v1alpha1/projects/{project_id}/locations/{region}'
        self.project_id = project_id
        self.project_number = project_number
        self.region = region
        
    def wait_for_completion(self, operation_id, message, sleep_time):
        Waits for a completion of a long running operation.
        
        api_url = f'{self.ann_parent}/operations/{operation_id}'
        start_time = datetime.datetime.utcnow()
        while True:
            response = self.authed_session.get(api_url)
            if response.status_code != 200:
                raise RuntimeError(response.json())
            if 'done' in response.json().keys():
                logging.info('Operation completed!')
                break
            elapsed_time = datetime.datetime.utcnow() - start_time
            logging.info('{}. Elapsed time since start: {}.'.format(
                message, str(elapsed_time)))
            time.sleep(sleep_time)
    
        return response.json()['response']
class IndexClient(ANNClient):
    Encapsulates a subset of control plane APIs 
    that manage ANN indexes.
    def __init__(self, project_id, project_number, region):
        super().__init__(project_id, project_number, region)
    def create_index(self, display_name, description, metadata):
        Creates an ANN Index.
    
        api_url = f'{self.ann_parent}/indexes'
    
        request_body = {
            'display_name': display_name,
            'description': description,
            'metadata': metadata
        }
    
        response = self.authed_session.post(api_url, data=json.dumps(request_body))
        if response.status_code != 200:
            raise RuntimeError(response.text)
        operation_id = response.json()['name'].split('/')[-1]
        
        return operation_id
    def list_indexes(self, display_name=None):
        Lists all indexes with a given display name or
        all indexes if the display_name is not provided.
    
        if display_name:
            api_url = f'{self.ann_parent}/indexes?filter=display_name="{display_name}"'
        else:
            api_url = f'{self.ann_parent}/indexes'
        response = self.authed_session.get(api_url).json()
        return response['indexes'] if response else []
    
    def delete_index(self, index_id):
        Deletes an ANN index.
        
        api_url = f'{self.ann_parent}/indexes/{index_id}'
        response = self.authed_session.delete(api_url)
        if response.status_code != 200:
            raise RuntimeError(response.text)
class IndexDeploymentClient(ANNClient):
    Encapsulates a subset of control plane APIs 
    that manage ANN endpoints and deployments.
    
    def __init__(self, project_id, project_number, region):
        super().__init__(project_id, project_number, region)
    def create_endpoint(self, display_name, vpc_name):
        Creates an ANN endpoint.
    
        api_url = f'{self.ann_parent}/indexEndpoints'
        network_name = f'projects/{self.project_number}/global/networks/{vpc_name}'
        request_body = {
            'display_name': display_name,
            'network': network_name
        }
        response = self.authed_session.post(api_url, data=json.dumps(request_body))
        if response.status_code != 200:
            raise RuntimeError(response.text)
        operation_id = response.json()['name'].split('/')[-1]
    
        return operation_id
    
    def list_endpoints(self, display_name=None):
        Lists all ANN endpoints with a given display name or
        all endpoints in the project if the display_name is not provided.
        
        if display_name:
            api_url = f'{self.ann_parent}/indexEndpoints?filter=display_name="{display_name}"'
        else:
            api_url = f'{self.ann_parent}/indexEndpoints'
        response = self.authed_session.get(api_url).json()
 
        return response['indexEndpoints'] if response else []
    
    def delete_endpoint(self, endpoint_id):
        Deletes an ANN endpoint.
        
        api_url = f'{self.ann_parent}/indexEndpoints/{endpoint_id}'
        
        response = self.authed_session.delete(api_url)
        if response.status_code != 200:
            raise RuntimeError(response.text)
        
        return response.json()
    
    def create_deployment(self, display_name, deployment_id, endpoint_id, index_id):
        Deploys an ANN index to an endpoint.
    
        api_url = f'{self.ann_parent}/indexEndpoints/{endpoint_id}:deployIndex'
        index_name = f'projects/{self.project_number}/locations/{self.region}/indexes/{index_id}'
        request_body = {
            'deployed_index': {
                'id': deployment_id,
                'index': index_name,
                'display_name': display_name
            }
        }
        response = self.authed_session.post(api_url, data=json.dumps(request_body))
        if response.status_code != 200:
            raise RuntimeError(response.text)
        operation_id = response.json()['name'].split('/')[-1]
        
        return operation_id
    
    def get_deployment_grpc_ip(self, endpoint_id, deployment_id):
        Returns a private IP address for a gRPC interface to 
        an Index deployment.
  
        api_url = f'{self.ann_parent}/indexEndpoints/{endpoint_id}'
        response = self.authed_session.get(api_url)
        if response.status_code != 200:
            raise RuntimeError(response.text)
            
        endpoint_ip = None
        if 'deployedIndexes' in response.json().keys():
            for deployment in response.json()['deployedIndexes']:
                if deployment['id'] == deployment_id:
                    endpoint_ip = deployment['privateEndpoints']['matchGrpcAddress']
                    
        return endpoint_ip
    
    def delete_deployment(self, endpoint_id, deployment_id):
        Undeployes an index from an endpoint.
        
        api_url = f'{self.ann_parent}/indexEndpoints/{endpoint_id}:undeployIndex'
        
        request_body = {
            'deployed_index_id': deployment_id
        }
    
        response = self.authed_session.post(api_url, data=json.dumps(request_body))
        if response.status_code != 200:
            raise RuntimeError(response.text)
        
        return response
    
index_client = IndexClient(PROJECT_ID, PROJECT_NUMBER, REGION)
deployment_client = IndexDeploymentClient(PROJECT_ID, PROJECT_NUMBER, REGION)
Explanation: Creating an ANN index deployment
Deploying an ANN index is a 3 step process:
1. Creating an index from source files
2. Creating an endpoint to access the index
3. Deploying the index to the endpoint
You will use the REST interface to invoke the AI Platform ANN Service Control Plane API that manages indexes, endpoints, and deployments.
After the index has been deployed you can submit matching requests using Online Querying API. In the experimental stage this API is only accessible through the gRPC interface.
Define helper classes to encapsulate the ANN Service REST API.
Currently, there is no Python client that encapsulates the ANN Service Control Plane API. The below code snippet defines a simple wrapper that encapsulates a subset of REST APIs used in this notebook.
End of explanation
indexes = index_client.list_indexes()
if not indexes:
    print('There are not any indexes registered with the service')
for index in indexes:
    print(index['name'])
Explanation: Create an ANN index
List all indexes registered with the ANN service
End of explanation
index_display_name = 'Song embeddings'
index_description = 'Song embeddings created BQML Matrix Factorization model'
index_metadata = {
    'contents_delta_uri': DATA_LOCATION,
    'config': {
        'dimensions': 50,
        'approximate_neighbors_count': 50,
        'distance_measure_type': 'DOT_PRODUCT_DISTANCE',
        'feature_norm_type': 'UNIT_L2_NORM',
        'tree_ah_config': {
            'child_node_count': 1000,
            'max_leaves_to_search': 100
         }
    }
}
logging.getLogger().setLevel(logging.INFO)
operation_id = index_client.create_index(index_display_name, 
                                          index_description,
                                          index_metadata)
response = index_client.wait_for_completion(operation_id, 'Creating index', 20)
print(response)
Explanation: Configure and create a new index based on the exported embeddings
Index creation is a long running operation. Be patient.
End of explanation
indexes = index_client.list_indexes(index_display_name)
for index in indexes:
    print(index['name'])
if indexes: 
    index_id = index['name'].split('/')[-1]
    print(f'Index: {index_id} will be used for deployment')
else:
    print('No indexes available for deployment')
Explanation: Verify that the index was created
End of explanation
endpoints = deployment_client.list_endpoints()
if not endpoints:
    print('There are not any endpoints registered with the service')
for endpoint in endpoints:
    print(endpoint['name'])
Explanation: Create the index deployment
List all endpoints registered with the ANN service
End of explanation
deployment_display_name = 'Song embeddings endpoint'
operation_id = deployment_client.create_endpoint(deployment_display_name, VPC_NAME)
response = index_client.wait_for_completion(operation_id, 'Waiting for endpoint', 10)
print(response)
Explanation: Create an index endpoint
End of explanation
endpoints = deployment_client.list_endpoints(deployment_display_name)
for endpoint in endpoints:
    print(endpoint['name'])
    
if endpoints: 
    endpoint_id = endpoint['name'].split('/')[-1]
    print(f'Endpoint: {endpoint_id} will be used for deployment')
else:
    print('No endpoints available for deployment')
Explanation: Verify that the endpoint was created
End of explanation
deployment_display_name = 'Song embeddings deployed index'
deployed_index_id = 'songs_embeddings_deployed_index'
Explanation: Deploy the index to the endpoint
Set the deployed index ID
The ID of the deployed index must be unique within your project.
End of explanation
response = index_client.wait_for_completion(operation_id, 'Waiting for deployment', 10)
operation_id = deployment_client.create_deployment(deployment_display_name, 
                                                   deployed_index_id,
                                                   endpoint_id,
                                                   index_id)
response = index_client.wait_for_completion(operation_id, 'Waiting for deployment', 10)
print(response)
Explanation: Deploy the index
Be patient. Index deployment is a long running operation
End of explanation
deployed_index_ip = deployment_client.get_deployment_grpc_ip(endpoint_id, deployed_index_id)
endpoint = f'{deployed_index_ip}:{MATCH_SERVICE_PORT}'
print(f'gRPC endpoint for the: {deployed_index_id} deployment is: {endpoint}')
Explanation: Querying the ANN service
You will use the gRPC interface to query the deployed index.
Retrieve the gRPC private endpoint for the ANN Match service
End of explanation
class MatchService(object):
    This is a wrapper around Online Querying gRPC interface.
    def __init__(self, endpoint, deployed_index_id):
        self.endpoint = endpoint
        self.deployed_index_id = deployed_index_id
    def single_match(
        self,
        embedding: List[float],
        num_neighbors: int) -> List[Tuple[str, float]]:
        Requests a match for a single embedding.
    
        match_request = match_pb2.MatchRequest(deployed_index_id=self.deployed_index_id,
                                               float_val=embedding,
                                               num_neighbors=num_neighbors)
        with grpc.insecure_channel(endpoint) as channel:
            stub = match_pb2_grpc.MatchServiceStub(channel)
            response = stub.Match(match_request)
    
        return [(neighbor.id, neighbor.distance) for neighbor in response.neighbor]
    def batch_match(
        self,
        embeddings: List[List[float]],
        num_neighbors: int) -> List[List[Tuple[str, float]]]:
        Requests matches ofr a list of embeddings.
    
        match_requests = [
            match_pb2.MatchRequest(deployed_index_id=self.deployed_index_id,
                                   float_val=embedding,
                                   num_neighbors=num_neighbors)
            for embedding in embeddings]
    
        batches_per_index = [
            match_pb2.BatchMatchRequest.BatchMatchRequestPerIndex(
                deployed_index_id=self.deployed_index_id,
                requests=match_requests)]
    
        batch_match_request = match_pb2.BatchMatchRequest(
            requests=batches_per_index)
    
        with grpc.insecure_channel(endpoint) as channel:
            stub = match_pb2_grpc.MatchServiceStub(channel)
            response = stub.BatchMatch(batch_match_request)
        
        matches = []
        for batch_per_index in response.responses:
            for match in batch_per_index.responses:
                matches.append(
                    [(neighbor.id, neighbor.distance) for neighbor in match.neighbor])
        
        return matches
    
match_service = MatchService(endpoint, deployed_index_id)
Explanation: Create a helper wrapper around the Match Service gRPC API.
The wrapper uses the pre-generated gRPC stub to the Online Querying gRPC interface.
End of explanation
%%bigquery df_embeddings
SELECT id, embedding
FROM `recommendations.item_embeddings` 
LIMIT 10
sample_embeddings = [list(embedding) for embedding in df_embeddings['embedding']]
sample_embeddings[0]
Explanation: Prepare sample data
Retrieve a few embeddings from the BigQuery embedding table
End of explanation
%%time 
single_match = match_service.single_match(sample_embeddings[0], 10)
single_match
Explanation: Run a single match query
The following call requests 10 closest neighbours for a single embedding.
End of explanation
%%time 
batch_match = match_service.batch_match(sample_embeddings[0:5], 3)
batch_match
Explanation: Run a batch match query
The following call requests 3 closest neighbours for each of the embeddings in a batch of 5.
End of explanation
for endpoint in deployment_client.list_endpoints():
    endpoint_id = endpoint['name'].split('/')[-1]
    if 'deployedIndexes' in endpoint.keys():
        for deployment in endpoint['deployedIndexes']:
            print('   Deleting index deployment: {} in the endpoint: {} '.format(deployment['id'], endpoint_id))
            deployment_client.delete_deployment(endpoint_id, deployment['id'])
    print('Deleting endpoint: {}'.format(endpoint['name']))
    deployment_client.delete_endpoint(endpoint_id)
Explanation: Clean up
WARNING
The below code will delete all ANN deployments, endpoints, and indexes in the configured project.
Delete index deployments and endpoints
End of explanation
for index in index_client.list_indexes():
    index_id = index['name'].split('/')[-1]
    print('Deleting index: {}'.format(index['name']))
    index_client.delete_index(index_id)
Explanation: Delete indexes
End of explanation
<END_TASK>