import os

import tensorflow as tf
import tensorflow_hub as hub

"""
Explanation: Custom TF-Hub Word Embedding with text2hub
Learning Objectives:
  1. Learn how to deploy AI Hub Kubeflow pipeline
  1. Learn how to configure the run parameters for text2hub
  1. Learn how to inspect text2hub generated artifacts and word embeddings in TensorBoard
  1. Learn how to run TF 1.x generated hub module in TF 2.0
Introduction
Pre-trained text embeddings such as TF-Hub modules are a great tool for building machine learning models for text features, since they capture relationships between words. These embeddings are generally trained on vast but generic text corpora like Wikipedia or Google News, which means that they are usually very good at representing generic text, but not so much when the text comes from a very specialized domain with unique vocabulary, such as in the medical field.
One problem in particular that arises when applying a TF-Hub text module which was pre-trained on a generic corpus to specialized text is that all of the unique, domain-specific words will be mapped to the same “out-of-vocabulary” (OOV) vector. By doing so we lose a very valuable part of the text information, because for specialized texts the most informative words are often the words that are very specific to that special domain. Another issue is that of commonly misspelled words from text gathered from say, customer feedback. Applying a generic pre-trained embedding will send the misspelled word to the OOV vectors, losing precious information. However, by creating a TF-Hub module tailored to the texts coming from that customer feedback means that common misspellings present in your real customer data will be part of the embedding vocabulary and should be close by closeby to the original word in the embedding space.
In this notebook, we will learn how to generate a text TF-hub module specific to a particular domain using the text2hub Kubeflow pipeline available on Google AI Hub. This pipeline takes as input a corpus of text stored in a GCS bucket and outputs a TF-Hub module to a GCS bucket. The generated TF-Hub module can then be reused both in TF 1.x or in TF 2.0 code by referencing the output GCS bucket path when loading the module. 
Our first order of business will be to learn how to deploy a Kubeflow pipeline, namely text2hub, stored in AI Hub to a Kubeflow cluster. Then we will dig into the pipeline run parameter configuration and review the artifacts produced by the pipeline during its run. These artifacts are meant to help you assess how good the domain specific TF-hub module you generated is. In particular, we will  explore the embedding space visually using TensorBoard projector, which provides a tool to list the nearest neighbors to a given word in the embedding space.
At last, we will explain how to run the generated module both in TF 1.x and TF 2.0. Running the module in TF 2.0 will necessite a small trick that’s useful to know in itself because it allows you to use all the TF 1.x modules in TF hub in TF 2.0 as a Keras layer. 
Each learning objective will correspond to a #TODO in this student lab notebook -- try to complete this notebook first and then review the solution notebook.
End of explanation
"""


PROJECT = "your-gcp-project-here" # REPLACE WITH YOUR PROJECT NAME
BUCKET = "your-gcp-bucket-here" # REPLACE WITH YOUR BUCKET NAME

os.environ["PROJECT"] = PROJECT
os.environ["BUCKET"] = BUCKET

!gsutil ls gs://$BUCKET || gsutil mb gs://$BUCKET

"""
Explanation: Replace by your GCP project and bucket:
End of explanation
"""


%%bash

URL=https://www.gutenberg.org/cache/epub/24564/pg24564.txt
OUTDIR=gs://$BUCKET/custom_embedding
CORPUS=surgery_manual.txt

curl $URL > $CORPUS
gsutil cp $CORPUS $OUTDIR/$CORPUS

"""
Explanation: Loading the dataset in GCS
The corpus we chose is one of Project Gutenberg medical texts: A Manual of the Operations of Surgery by Joseph Bell, containing very specialized language. 
The first thing to do is to upload the text into a GCS bucket:
End of explanation
"""


!echo gs://$BUCKET/custom_embedding/surgery_manual.txt

"""
Explanation: It has very specialized language such as 
On the surface of the abdomen the position of this vessel would be 
indicated by a line drawn from about an inch on either side of the 
umbilicus to the middle of the space between the symphysis pubis 
and the crest of the ilium.
Now let's go over the steps involved in creating your own embedding from that corpus.
Lab Task 1:
Complete Step 1 through Step 3 below to deploy an AI Hub Kubeflow pipeline.
Step 1: Download the text2hub pipeline from AI Hub
Go on AI Hub and search for the text2hub pipeline, or just follow this link.
You'll land onto a page describing text2hub. Click on the "Download" button on that page to download the Kubeflow pipeline and click Accept.

The text2hub pipeline is a KubeFlow pipeline that comprises three components; namely:


The text2cooc component that computes a word co-occurrence matrix
from a corpus of text


The cooc2emb component that factorizes the
co-occurrence matrix using Swivel into
the word embeddings exported as a tsv file


The emb2hub component that takes the word
embedding file and generates a TF Hub module from it


Each component is implemented as a Docker container image that's stored into Google Cloud Docker registry, gcr.io. The pipeline.tar.gz file that you downloaded is a yaml description of how these containers need to be composed as well as where to find the corresponding images. 
Remark: Each component can be run individually as a single component pipeline in exactly the same manner as the text2hub pipeline. On AI Hub, each component has a pipeline page describing it and from where you can download the associated single component pipeline:

text2cooc
cooc2emb
emb2hub

Step 2: Upload the pipeline to the Kubeflow cluster
Go to your Kubeflow cluster dashboard or navigate to Navigation menu &gt; AI Platform &gt; Pipelines and click Open Pipelines Dashboard then click on the Pipelines tab to create a new pipeline. You'll be prompted to upload the pipeline file you have just downloaded, click Upload Pipeline. Rename the generated pipeline name to be text2hub to keep things nice and clean.

Step 3: Create a pipeline run
After uploading the pipeline, you should see text2hub appear on the pipeline list. Click on it. This will bring you to a page describing the pipeline (explore!) and allowing you to create a run. You can inspect the input and output parameters of each of the pipeline components by clicking on the component node in the graph representing the pipeline. Click Create Run.

Lab Task 2:
Complete Step 4 below to configure the run parameters for text2hub.
Step 4: Enter the run parameters
text2hub has the following run parameters you can configure:
Argument                                         | Description                                                                           | Optional | Data Type | Accepted values | Default
------------------------------------------------ | ------------------------------------------------------------------------------------- | -------- | --------- | --------------- | -------
gcs-path-to-the-text-corpus                      | A Cloud Storage location pattern (i.e., glob) where the text corpus will be read from | False    | String    | gs://...        | -
gcs-directory-path-for-pipeline-output           | A Cloud Storage directory path where the pipeline output will be exported             | False    | String    | gs://...        | -
number-of-epochs                                 | Number of epochs to train the embedding algorithm (Swivel) on                         | True     | Integer   | -               | 40
embedding-dimension                              | Number of components of the generated embedding vectors                               | True     | Integer   | -               | 128
co-occurrence-word-window-size                   | Size of the sliding word window where co-occurrences are extracted from               | True     | Integer   | -               | 10
number-of-out-of-vocabulary-buckets              | Number of out-of-vocabulary buckets                                                   | True     | Integer   | -               | 1
minimum-occurrences-for-a-token-to-be-considered | Minimum number of occurrences for a token to be included in the vocabulary            | True     | Integer   | -               | 5
You can leave most parameters with their default values except for
gcs-path-to-the-test-corpus whose value should be set to
End of explanation
"""


!echo gs://$BUCKET/custom_embedding

"""
Explanation: and for gcs-directory-path-for-pipeline-output which we will set to
End of explanation
"""


!echo tensorboard --port 8080 --logdir gs://$BUCKET/custom_embedding/embeddings

"""
Explanation: Remark: gcs-path-to-the-test-corpus will accept a GCS pattern like gs://BUCKET/data/*.txt or simply a path like gs://BUCKET/data/ to a GCS directory. All the files that match the pattern or that are in that directory will be parsed to create the word embedding TF-Hub module. 

Make sure to choose experiment default. Once these values have been set, you can start the run by clicking on Start.
Lab Task 3:
Complete Step 5 below to use Tensorboard to inspect the text2hub generated artifacts and word embeddings.
Step 5: Inspect the run artifacts
Once the run has started you can see its state by going to the Experiments tab and clicking on the name of the run (here "text2hub-1"). 

It will show you the pipeline graph. The components in green have successfuly completed. You can then click on them and look at the artifacts that these components have produced.
The text2cooc components has "co-occurrence extraction summary" showing you the GCS path where the co-occurrence data has been saved. Their is a corresponding link that you can paste into your browser to inspect the co-occurrence data from the GCS browser. Some statistics about the vocabulary are given to you such as the most and least frequent tokens. You can also download the vocabulary file containing the token to be embedded. 

The cooc2emb has three artifacts
* An "Embedding Extraction Summary" providing the information as where the model chekpoints and the embedding tables are exported to on GCP
* A similarity matrix from a random sample of words giving you an indication whether the model associates close-by vectors to similar words
* An button to start TensorBoard from the UI to inspect the model and visualize the word embeddings

We can have a look at the word embedding visualization provided by TensorBoard. Select the TF version: TensorFlow 1.14.0. Start TensorBoard by clicking on Start Tensorboard and then Open Tensorboard buttons, and then select "Projector".
Remark: The projector tab may take some time to appear. If it takes too long it may be that your Kubeflow cluster is running an incompatible version of TensorBoard (your TB version should be between 1.13 and 1.15). If that's the case, just run Tensorboard from CloudShell or locally by issuing the following command.
End of explanation
"""


MODULE = "gs://{bucket}/custom_embedding/hub-module".format(bucket=BUCKET)
MODULE

"""
Explanation: The projector view will present you with a representation of the word vectors in a 3 dimensional space (the dim is reduced through PCA) that you can interact with. Enter in the search tool a few words like "ilium" and points in the 3D space will light up. 

If you click on a word vector, you'll see appear the n nearest neighbors of that word in the embedding space. The nearset neighbors are both visualized in the center panel and presented as a flat list on the right. 
Explore the nearest neighbors of a given word and see if they kind of make sense. This will give you a rough understanding of the embedding quality. If it nearest neighbors do not make sense after trying for a few key words, you may need rerun text2hub, but this time with either more epochs or more data. Reducing the embedding dimension may help as well as modifying the co-occurence window size (choose a size that make sense given how your corpus is split into lines.)

The emb2hub artifacts give you a snippet of TensorFlow 1.x code showing you how to re-use the generated TF-Hub module in your code. We will demonstrate how to use the TF-Hub module in TF 2.0 in the next section.

Step 7: Using the generated TF-Hub module (TODO)
Let's see now how to load the TF-Hub module generated by text2hub in TF 2.0.
We first store the GCS bucket path where the TF-Hub module has been exported into a variable:
End of explanation
"""


med_embed = # TODO: Your code goes here.

"""
Explanation: Now we are ready to create a KerasLayer out of our custom text embedding.
Lab Task 4a:
Complete the code below to create a KerasLayer using the MODULE defined above.
End of explanation
"""


outputs = # TODO: Your code goes here.
outputs

"""
Explanation: That layer when called with a list of sentences will create a sentence vector for each sentence by averaging the word vectors of the sentence.
Lab Task 4b:
Use the med_embed layer you created above to find text embeddings for the words/phrases 'ilium', 'I have a fracture', and 'aneurism'.
End of explanation
"""


