#!/usr/bin/env python
# coding: utf-8

#
# # Assignment 2: Transformer Summarizer
#
# Welcome to the second assignment of course 4. In this assignment you will explore summarization using the transformer model. Yes, you will implement the transformer decoder from scratch, but we will slowly walk you through it. There are many hints in this notebook so feel free to use them as needed.
#
#
#
# ## Important Note on Submission to the AutoGrader
#
# Before submitting your assignment to the AutoGrader, please make sure you are not doing the following:
#
# 1. You have not added any _extra_ `print` statement(s) in the assignment.
# 2. You have not added any _extra_ code cell(s) in the assignment.
# 3. You have not changed any of the function parameters.
# 4. You are not using any global variables inside your graded exercises. Unless specifically instructed to do so, please refrain from it and use the local variables instead.
# 5. You are not changing the assignment code where it is not required, like creating _extra_ variables.
#
# If you do any of the following, you will get something like, `Grader not found` (or similarly unexpected) error upon submitting your assignment. Before asking for help/debugging the errors in your assignment, check for these first. If this is the case, and you don't remember the changes you have made, you can get a fresh copy of the assignment by following these [instructions](https://www.coursera.org/learn/attention-models-in-nlp/supplement/shBOS/how-to-refresh-your-workspace).

# ## Outline
#
# - [Introduction](#0)
# - [Part 1: Importing the dataset](#1)
#     - [1.1 Encode & Decode helper functions](#1.1)
#     - [1.2 Defining parameters](#1.2)
#     - [1.3 Exploring the data](#1.3)
# - [Part 2: Summarization with transformer](#2)
#     - [2.1 Dot product attention](#2.1)
#         - [Exercise 01](#ex01)
#     - [2.2 Causal Attention](#2.2)
#         - [Exercise 02](#ex02)
#     - [2.3 Transformer decoder block](#2.3)
#         - [Exercise 03](#ex03)
#     - [2.4 Transformer Language model](#2.4)
#         - [Exercise 04](#ex04)
# - [Part 3: Training](#3)
#     - [3.1 Training the model](#3.1)
#         - [Exercise 05](#ex05)
# - [Part 4: Evaluation](#4)
#     - [4.1 Loading in a trained model](#4.1)
# - [Part 5: Testing with your own input](#5)
#     - [Exercise 6](#ex06)
#     - [5.1 Greedy decoding](#5.1)
#         - [Exercise 07](#ex07)

#
# ### Introduction
#
# Summarization is an important task in natural language processing and could be useful for a consumer enterprise. For example, bots can be used to scrape articles, summarize them, and then you can use sentiment analysis to identify the sentiment about certain stocks. Anyways who wants to read an article or a long email today, when you can build a transformer to summarize text for you. Let's get started, by completing this assignment you will learn to:
#
# - Use built-in functions to preprocess your data
# - Implement DotProductAttention
# - Implement Causal Attention
# - Understand how attention works
# - Build the transformer model
# - Evaluate your model
# - Summarize an article
#
# As you can tell, this model is slightly different than the ones you have already implemented. This is heavily based on attention and does not rely on sequences, which allows for parallel computing.

# In[1]:

import sys
import os
import w2_tests
import numpy as np

import textwrap
wrapper = textwrap.TextWrapper(width=70)

import trax
from trax import layers as tl
from trax.fastmath import numpy as jnp

# to print the entire np array
np.set_printoptions(threshold=sys.maxsize)

# In[2]:

#
# ## Part 1: Importing the dataset

# Trax makes it easy to work with Tensorflow's datasets:

# In[3]:

# This will download the dataset if no data_dir is specified.
# so we have the data already in 'data/' for you

# Importing CNN/DailyMail articles dataset
train_stream_fn = trax.data.TFDS('cnn_dailymail',
                                 data_dir='data/',
                                 keys=('article', 'highlights'),
                                 train=True)

# This should be much faster as the data is downloaded already.
eval_stream_fn = trax.data.TFDS('cnn_dailymail',
                                data_dir='data/',
                                keys=('article', 'highlights'),
                                train=False)

#
# ## 1.1 Tokenize & Detokenize helper functions
#
# Just like in the previous assignment, the cell above loads in the encoder for you. Given any data set, you have to be able to map words to their indices, and indices to their words. The inputs and outputs to your [Trax](https://github.com/google/trax) models are usually tensors of numbers where each number corresponds to a word. If you were to process your data manually, you would have to make use of the following:
#
# -  word2Ind:  a dictionary mapping the word to its index.
# -  ind2Word: a dictionary mapping the index to its word.
# -  word2Count: a dictionary mapping the word to the number of times it appears.
# -  num_words: total number of words that have appeared.
#
# Since you have already implemented these in previous assignments of the specialization, we will provide you with helper functions that will do this for you. Run the cell below to get the following functions:
#
# -  tokenize:  converts a text sentence to its corresponding token list (i.e. list of indices). Also converts words to subwords.
# -  detokenize:  converts a token list to its corresponding sentence (i.e. string).

# In[4]:

def tokenize(input_str, EOS=1):
    """Input str to features dict, ready for inference"""

    # Use the trax.data.tokenize method. It takes streams and returns streams,
    # we get around it by making a 1-element stream with `iter`.
    inputs =  next(trax.data.tokenize(iter([input_str]),
                                      vocab_dir='vocab_dir/',
                                      vocab_file='summarize32k.subword.subwords'))

    # Mark the end of the sentence with EOS
    return list(inputs) + [EOS]

def detokenize(integers):
    """List of ints to str"""

    s = trax.data.detokenize(integers,
                             vocab_dir='vocab_dir/',
                             vocab_file='summarize32k.subword.subwords')

    return wrapper.fill(s)

#
#
# ## 1.2 Preprocessing for Language Models: Concatenate It!
#
# This week you will use a language model -- Transformer Decoder -- to solve
# an input-output problem. As you know, language models only predict the next
# word, they have no notion of inputs. To create a single input suitable for
# a language model, we concatenate inputs with targets putting a separator
# in between. We also need to create a mask -- with 0s at inputs and 1s at targets -- so that the model is not penalized for mis-predicting the article and only focuses on the summary. See the preprocess function below for how this is done.

# In[5]:

# Special tokens
SEP = 0  # Padding or separator token
EOS = 1  # End of sentence token

# Concatenate tokenized inputs and targets using 0 as separator.
def preprocess(stream):
    for (article, summary) in stream:
        joint = np.array(list(article) + [EOS, SEP] + list(summary) + [EOS])
        mask = [0] * (len(list(article)) + 2) + [1] * (len(list(summary)) + 1) # Accounting for EOS and SEP
        yield joint, joint, np.array(mask)

# You can combine a few data preprocessing steps into a pipeline like this.
input_pipeline = trax.data.Serial(
    # Tokenizes
    trax.data.Tokenize(vocab_dir='vocab_dir/',
                       vocab_file='summarize32k.subword.subwords'),
    # Uses function defined above
    preprocess,
    # Filters out examples longer than 2048
    trax.data.FilterByLength(2048)
)

# Apply preprocessing to data streams.
train_stream = input_pipeline(train_stream_fn())
eval_stream = input_pipeline(eval_stream_fn())

train_input, train_target, train_mask = next(train_stream)

assert sum((train_input - train_target)**2) == 0  # They are the same in Language Model (LM).

# In[6]:

# prints mask, 0s on article, 1s on summary
print(f'Single example mask:\n\n {train_mask}')

# In[7]:

# prints: [Example][<EOS>][<pad>][Example Summary][<EOS>]
print(f'Single example:\n\n {detokenize(train_input)}')

#
#
# ## 1.3 Batching with bucketing
#
# As in the previous week, we use bucketing to create batches of data.

# In[8]:

# Bucketing to create batched generators.

# Buckets are defined in terms of boundaries and batch sizes.
# Batch_sizes[i] determines the batch size for items with length < boundaries[i]
# So below, we'll take a batch of 16 sentences of length < 128 , 8 of length < 256,
# 4 of length < 512. And so on.
boundaries =  [128, 256,  512, 1024]
batch_sizes = [16,    8,    4,    2, 1]

# Create the streams.
train_batch_stream = trax.data.BucketByLength(
    boundaries, batch_sizes)(train_stream)

eval_batch_stream = trax.data.BucketByLength(
    boundaries, batch_sizes)(eval_stream)

# In[9]:

# Every execution will result in generation of a different article
# Try running this cell multiple times to see how the length of the examples affects the batch size
input_batch, _, mask_batch = next(train_batch_stream)

# Shape of the input_batch
input_batch.shape

# In[10]:

# print corresponding integer values
print(input_batch[0])

# Things to notice:
#  - First we see the corresponding values of the words.
#  - The first 1, which represents the `<EOS>` tag of the article.
#  - Followed by a 0, which represents a `<pad>` tag.
#  - After the first 0 (`<pad>` tag) the corresponding values are of the words that are used for the summary of the article.
#  - The second 1 represents the `<EOS>` tag for the summary.
#  - All the trailing 0s represent `<pad>` tags which are appended to maintain consistent length (If you don't see them then it would mean it is already of max length)
#

# In[11]:

# print the article and its summary
print('Article:\n\n', detokenize(input_batch[0]))

# You can see that the data has the following structure:
# -  [Article]  -> `<EOS>` -> `<pad>` ->  [Article Summary]  -> `<EOS>` -> (possibly) multiple `<pad>`
#
# The loss is taken only on the summary using cross_entropy as loss function.

#
# # Part 2: Summarization with transformer
#
# Now that we have given you the data generator and have handled the preprocessing for you, it is time for you to build your own model. We saved you some time because we know you have already preprocessed data before in this specialization, so we would rather you spend your time doing the next steps.
#
# You will be implementing the attention from scratch and then using it in your transformer model. Concretely, you will understand how attention works, how you use it to connect the encoder and the decoder.
#
#
#
#
# ## 2.1 Dot product attention
#
# Now you will implement dot product attention which takes in a query, key, value, and a mask. It returns the output.
#
#
#
#
# Here are some helper functions that will help you create tensors and display useful information:
#    - `create_tensor`  creates a `jax numpy array` from a list of lists.
#    - `display_tensor` prints out the shape and the actual tensor.

# In[12]:

def create_tensor(t):
    """Create tensor from list of lists"""
    return jnp.array(t)

def display_tensor(t, name):
    """Display shape and tensor"""
    print(f'{name} shape: {t.shape}\n')
    print(f'{t}\n')

# Before implementing it yourself, you can play around with a toy example of `dot product attention` without the softmax  operation. Technically it would not be `dot product attention` without the softmax but this is done to avoid giving away too much of the answer and the idea is to display these tensors to give you a sense of how they look like.
#
# The formula for attention is this one:
#
# $$
# \text { Attention }(Q, K, V)=\operatorname{softmax}\left(\frac{Q K^{T}}{\sqrt{d_{k}}}+{M}\right) V\tag{1}\
# $$
#
# $d_{k}$ stands for the dimension of queries and keys.
#
# The `query`, `key`, `value` and `mask` vectors are provided for this example.
#
# Notice that the masking is done using very negative values that will yield a similar effect to using $-\infty $.

# In[13]:

q = create_tensor([[1, 0, 0], [0, 1, 0]])
display_tensor(q, 'query')
k = create_tensor([[1, 2, 3], [4, 5, 6]])
display_tensor(k, 'key')
v = create_tensor([[0, 1, 0], [1, 0, 1]])
display_tensor(v, 'value')
m = create_tensor([[0, 0], [-1e9, 0]])
display_tensor(m, 'mask')

# **Expected Output:**
# ```CPP
# query shape: (2, 3)
#
# [[1 0 0]
#  [0 1 0]]
#
# key shape: (2, 3)
#
# [[1 2 3]
#  [4 5 6]]
#
# value shape: (2, 3)
#
# [[0 1 0]
#  [1 0 1]]
#
# mask shape: (2, 2)
#
# [[ 0.e+00  0.e+00]
#  [-1.e+09  0.e+00]]
#
# ```

# In[14]:

q_dot_k = q @ k.T / jnp.sqrt(3)
display_tensor(q_dot_k, 'query dot key')

# **Expected Output:**
# ```CPP
# query dot key shape: (2, 2)
#
# [[0.57735026 2.309401  ]
#  [1.1547005  2.8867514 ]]
# ```

# In[15]:

masked = q_dot_k + m
display_tensor(masked, 'masked query dot key')

# **Expected Output:**
# ```CPP
# masked query dot key shape: (2, 2)
#
# [[ 5.7735026e-01  2.3094010e+00]
#  [-1.0000000e+09  2.8867514e+00]]
# ```

# In[16]:

display_tensor(masked @ v, 'masked query dot key dot value')

# **Expected Output:**
# ```CPP
# masked query dot key dot value shape: (2, 3)
#
# [[ 2.3094010e+00  5.7735026e-01  2.3094010e+00]
#  [ 2.8867514e+00 -1.0000000e+09  2.8867514e+00]]
# ```

# In order to use the previous dummy tensors to test some of the graded functions, a batch dimension should be added to them so they mimic the shape of real-life examples. The mask is also replaced by a version of it that resembles the one that is used by trax:

# In[17]:

q_with_batch = q[None,:]
display_tensor(q_with_batch, 'query with batch dim')
k_with_batch = k[None,:]
display_tensor(k_with_batch, 'key with batch dim')
v_with_batch = v[None,:]
display_tensor(v_with_batch, 'value with batch dim')
m_bool = create_tensor([[True, True], [False, True]])
display_tensor(m_bool, 'boolean mask')

# **Expected Output:**
# ```CPP
# query with batch dim shape: (1, 2, 3)
#
# [[[1 0 0]
#   [0 1 0]]]
#
# key with batch dim shape: (1, 2, 3)
#
# [[[1 2 3]
#   [4 5 6]]]
#
# value with batch dim shape: (1, 2, 3)
#
# [[[0 1 0]
#   [1 0 1]]]
#
# boolean mask shape: (2, 2)
#
# [[ True  True]
#  [False  True]]
# ```

#
# ### Exercise 01
#
# **Instructions:** Implement the dot product attention. Concretely, implement the following equation
#
#
# $$
# \text { Attention }(Q, K, V)=\operatorname{softmax}\left(\frac{Q K^{T}}{\sqrt{d_{k}}}+{M}\right) V\tag{1}\
# $$
#
# $Q$ - query,
# $K$ - key,
# $V$ - values,
# $M$ - mask,
# ${d_k}$ - depth/dimension of the queries and keys (used for scaling down)
#
# You can implement this formula either by `trax` numpy (trax.math.numpy) or regular `numpy` but it is recommended to use `jnp`.
#
# Something to take into consideration is that within trax, the masks are tensors of `True/False` values not 0's and $-\infty$ as in the previous example. Within the graded function don't think of applying the mask by summing up matrices, instead use `jnp.where()` and treat the **mask as a tensor of boolean values with `False` for values that need to be masked and True for the ones that don't.**
#
# Also take into account that the real tensors are far more complex than the toy ones you just played with. Because of this avoid using shortened operations such as `@` for dot product or `.T` for transposing. Use `jnp.matmul()` and `jnp.swapaxes()` instead.
#
# This is the self-attention block for the transformer decoder. Good luck!

# In[18]:

# UNQ_C1
# GRADED FUNCTION: DotProductAttention
def DotProductAttention(query, key, value, mask):
    """Dot product self-attention.
    Args:
        query (jax.interpreters.xla.DeviceArray): array of query representations with shape (L_q by d)
        key (jax.interpreters.xla.DeviceArray): array of key representations with shape (L_k by d)
        value (jax.interpreters.xla.DeviceArray): array of value representations with shape (L_k by d) where L_v = L_k
        mask (jax.interpreters.xla.DeviceArray): attention-mask, gates attention with shape (L_q by L_k)

    Returns:
        jax.interpreters.xla.DeviceArray: Self-attention array for q, k, v arrays. (L_q by L_k)
    """

    assert query.shape[-1] == key.shape[-1] == value.shape[-1], "Embedding dimensions of q, k, v aren't all the same"

    ### START CODE HERE (REPLACE INSTANCES OF 'None' WITH YOUR CODE) ###
    # Save depth/dimension of the query embedding for scaling down the dot product
    depth = query.shape[-1]

    # Calculate scaled query key dot product according to formula above
    dots = jnp.matmul(query, jnp.swapaxes(key, -1, -2)) / jnp.sqrt(depth)

    # Apply the mask
    if mask is not None: # You do not need to replace the 'None' on this line
        dots = jnp.where(mask, dots, jnp.full_like(dots, -1e9))

    # Softmax formula implementation
    # Use trax.fastmath.logsumexp of masked_qkT to avoid underflow by division by large numbers
    # Note: softmax = None
    logsumexp = trax.fastmath.logsumexp(dots, axis=-1, keepdims=True)

    # Take exponential of dots minus logsumexp to get softmax
    # Use jnp.exp()
    dots = jnp.exp(dots - logsumexp)

    # Multiply dots by value to get self-attention
    # Use jnp.matmul()
    attention = jnp.matmul(dots, value)

    ## END CODE HERE ###

    return attention

# In[19]:

DotProductAttention(q_with_batch, k_with_batch, v_with_batch, m_bool)

# **Expected Output:**
# ```CPP
# DeviceArray([[[0.8496746 , 0.15032545, 0.8496746 ],
#               [1.        , 0.        , 1.        ]]], dtype=float32)
# ```

# In[20]:

# UNIT TEST
# test DotProductAttention
w2_tests.test_DotProductAttention(DotProductAttention)

#
#
# ## 2.2 Causal Attention
#
# Now you are going to implement causal attention: multi-headed attention with a mask to attend only to words that occurred before.
#
#
#
# In the image above, a word can see everything that is before it, but not what is after it. To implement causal attention, you will have to transform vectors and do many reshapes. You will need to implement the functions below.
#
#
#
# ### Exercise 02
#
# Implement the following functions that will be needed for Causal Attention:
#
# -  compute_attention_heads : Gets an input $x$ of dimension (n_batch, seqlen, n_heads $\times$ d_head) and splits the last (depth) dimension and stacks it to the zeroth dimension to allow matrix multiplication (n_batch $\times$ n_heads, seqlen, d_head).
# -  dot_product_self_attention : Creates a mask matrix with `False` values above the diagonal and `True` values below and calls DotProductAttention which implements dot product self attention.
# -  compute_attention_output : Undoes compute_attention_heads by splitting first (vertical) dimension and stacking in the last (depth) dimension (n_batch, seqlen, n_heads $\times$ d_head). These operations concatenate (stack/merge) the heads.
#
# Next there are some toy tensors which may serve to give you an idea of the data shapes and opperations involved in Causal Attention. They are also useful to test out your functions!

# In[21]:

tensor2d = create_tensor(q)
display_tensor(tensor2d, 'query matrix (2D tensor)')

tensor4d2b = create_tensor([[q, q], [q, q]])
display_tensor(tensor4d2b, 'batch of two (multi-head) collections of query matrices (4D tensor)')

tensor3dc = create_tensor([jnp.concatenate([q, q], axis = -1)])
display_tensor(tensor3dc, 'one batch of concatenated heads of query matrices (3d tensor)')

tensor3dc3b = create_tensor([jnp.concatenate([q, q], axis = -1), jnp.concatenate([q, q], axis = -1), jnp.concatenate([q, q], axis = -1)])
display_tensor(tensor3dc3b, 'three batches of concatenated heads of query matrices (3d tensor)')

# It is important to know that the following 3 functions would normally be defined within the `CausalAttention` function further below.
#
# However this makes these functions harder to test. Because of this, these functions are shown individually using a `closure` (when necessary) that simulates them being inside of the `CausalAttention` function. This is done because they rely on some variables that can be accessed from within `CausalAttention`.
#
# ### Support Functions
#
#  compute_attention_heads : Gets an input $x$ of dimension (n_batch, seqlen, n_heads $\times$ d_head) and splits the last (depth) dimension and stacks it to the zeroth dimension to allow matrix multiplication (n_batch $\times$ n_heads, seqlen, d_head).
#
# **For the closures you only have to fill the inner function.**

# In[22]:

# UNQ_C2
# GRADED FUNCTION: compute_attention_heads_closure
def compute_attention_heads_closure(n_heads, d_head):
    """ Function that simulates environment inside CausalAttention function.
    Args:
        d_head (int):  dimensionality of heads
        n_heads (int): number of attention heads
    Returns:
        function: compute_attention_heads function
    """

    def compute_attention_heads(x):
        """ Compute the attention heads.
        Args:
            x (jax.interpreters.xla.DeviceArray): tensor with shape (n_batch, seqlen, n_heads X d_head).
        Returns:
            jax.interpreters.xla.DeviceArray: reshaped tensor with shape (n_batch X n_heads, seqlen, d_head).
        """
        ### START CODE HERE ###
        # (REPLACE INSTANCES OF 'None' WITH YOUR CODE)

        # Size of the x's batch dimension
        batch_size = x.shape[0]
        # Length of the sequence
        # Should be size of x's first dimension without counting the batch dim
        seqlen = x.shape[1]
        # Reshape x using jnp.reshape()
        # n_batch, seqlen, n_heads*d_head -> n_batch, seqlen, n_heads, d_head
        x = jnp.reshape(x, (batch_size, seqlen, n_heads, d_head))
        # Transpose x using jnp.transpose()
        # n_batch, seqlen, n_heads, d_head -> n_batch, n_heads, seqlen, d_head
        # Note that the values within the tuple are the indexes of the dimensions of x and you must rearrange them
        x = jnp.transpose(x, (0, 2, 1, 3))
        # Reshape x using jnp.reshape()
        # n_batch, n_heads, seqlen, d_head -> n_batch*n_heads, seqlen, d_head
        x = jnp.reshape(x, (-1, seqlen, d_head))

        ### END CODE HERE ###

        return x
    return compute_attention_heads

# In[23]:

display_tensor(tensor3dc3b, "input tensor")
result_cah = compute_attention_heads_closure(2,3)(tensor3dc3b)
display_tensor(result_cah, "output tensor")

# **Expected Output:**
# ```CPP
# input tensor shape: (3, 2, 6)
#
# [[[1 0 0 1 0 0]
#   [0 1 0 0 1 0]]
#
#  [[1 0 0 1 0 0]
#   [0 1 0 0 1 0]]
#
#  [[1 0 0 1 0 0]
#   [0 1 0 0 1 0]]]
#
# output tensor shape: (6, 2, 3)
#
# [[[1 0 0]
#   [0 1 0]]
#
#  [[1 0 0]
#   [0 1 0]]
#
#  [[1 0 0]
#   [0 1 0]]
#
#  [[1 0 0]
#   [0 1 0]]
#
#  [[1 0 0]
#   [0 1 0]]
#
#  [[1 0 0]
#   [0 1 0]]]
# ```

# In[24]:

# UNIT TEST
# test compute_attention_heads_closure
w2_tests.test_compute_attention_heads_closure(compute_attention_heads_closure)

#  dot_product_self_attention : Creates a mask matrix with `False` values above the diagonal and `True` values below and calls DotProductAttention which implements dot product self attention.

# In[25]:

# UNQ_C3
# GRADED FUNCTION: dot_product_self_attention
def dot_product_self_attention(q, k, v):
    """ Masked dot product self attention.
    Args:
        q (jax.interpreters.xla.DeviceArray): queries.
        k (jax.interpreters.xla.DeviceArray): keys.
        v (jax.interpreters.xla.DeviceArray): values.
    Returns:
        jax.interpreters.xla.DeviceArray: masked dot product self attention tensor.
    """
    ### START CODE HERE (REPLACE INSTANCES OF 'None' with your code) ###

    # Hint: mask size should be equal to L_q. Remember that q has shape (batch_size, L_q, d)
    mask_size = q.shape[-2]

    # Creates a matrix with ones below the diagonal and 0s above. It should have shape (1, mask_size, mask_size)
    # Notice that 1's and 0's get casted to True/False by setting dtype to jnp.bool_
    # Use jnp.tril() - Lower triangle of an array and jnp.ones()
    mask = jnp.tril(jnp.ones((1, mask_size, mask_size), dtype=jnp.bool_), k=0)

    ### END CODE HERE ###

    return DotProductAttention(q, k, v, mask)

# In[26]:

dot_product_self_attention(q_with_batch, k_with_batch, v_with_batch)

# **Expected Output:**
# ```CPP
# DeviceArray([[[0.        , 1.        , 0.        ],
#               [0.8496746 , 0.15032543, 0.8496746 ]]], dtype=float32)
# ```

# In[27]:

# UNIT TEST
# test dot_product_self_attention
w2_tests.test_dot_product_self_attention(dot_product_self_attention)

#  compute_attention_output : Undoes compute_attention_heads by splitting first (vertical) dimension and stacking in the last (depth) dimension (n_batch, seqlen, n_heads $\times$ d_head). These operations concatenate (stack/merge) the heads.

# In[28]:

# UNQ_C4
# GRADED FUNCTION: compute_attention_output_closure
def compute_attention_output_closure(n_heads, d_head):
    """ Function that simulates environment inside CausalAttention function.
    Args:
        d_head (int):  dimensionality of heads
        n_heads (int): number of attention heads
    Returns:
        function: compute_attention_output function
    """

    def compute_attention_output(x):
        """ Compute the attention output.
        Args:
            x (jax.interpreters.xla.DeviceArray): tensor with shape (n_batch X n_heads, seqlen, d_head).
        Returns:
            jax.interpreters.xla.DeviceArray: reshaped tensor with shape (n_batch, seqlen, n_heads X d_head).
        """
        ### START CODE HERE (REPLACE INSTANCES OF 'None' WITH YOUR CODE) ###

        # Length of the sequence
        # Should be size of x's first dimension without counting the batch dim
        seqlen = x.shape[1]
        # Reshape x using jnp.reshape() to shape (n_batch, n_heads, seqlen, d_head)
        x = jnp.reshape(x, ( -1, n_heads, seqlen, d_head))
        # Transpose x using jnp.transpose() to shape (n_batch, seqlen, n_heads, d_head)
        x = jnp.transpose(x, ( 0, 2, 1 , 3))

        ### END CODE HERE ###

        # Reshape to allow to concatenate the heads
        return jnp.reshape(x, (-1, seqlen, n_heads * d_head))
    return compute_attention_output

# In[29]:

display_tensor(result_cah, "input tensor")
result_cao = compute_attention_output_closure(2,3)(result_cah)
display_tensor(result_cao, "output tensor")

# **Expected Output:**
# ```CPP
# input tensor shape: (6, 2, 3)
#
# [[[1 0 0]
#   [0 1 0]]
#
#  [[1 0 0]
#   [0 1 0]]
#
#  [[1 0 0]
#   [0 1 0]]
#
#  [[1 0 0]
#   [0 1 0]]
#
#  [[1 0 0]
#   [0 1 0]]
#
#  [[1 0 0]
#   [0 1 0]]]
#
# output tensor shape: (3, 2, 6)
#
# [[[1 0 0 1 0 0]
#   [0 1 0 0 1 0]]
#
#  [[1 0 0 1 0 0]
#   [0 1 0 0 1 0]]
#
#  [[1 0 0 1 0 0]
#   [0 1 0 0 1 0]]]
# ```

# In[30]:

# UNIT TEST
# test compute_attention_output_closure
w2_tests.test_compute_attention_output_closure(compute_attention_output_closure)

# ### Causal Attention Function
#
# Now it is time for you to put everything together within the `CausalAttention` or Masked multi-head attention function:

#
#
# **Instructions:** Implement the causal attention.
# Your model returns the causal attention through a $tl.Serial$ with the following:
#
# -  [tl.Branch](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.combinators.Branch) : consisting of 3 [tl.Dense(d_feature), ComputeAttentionHeads] to account for the queries, keys, and values.
# -  [tl.Fn](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.base.Fn): Takes in dot_product_self_attention function and uses it to compute the dot product using $Q$, $K$, $V$.
# -  [tl.Fn](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.base.Fn): Takes in compute_attention_output_closure to allow for parallel computing.
# -  [tl.Dense](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Dense): Final Dense layer, with dimension `d_feature`.
#
# Remember that in order for trax to properly handle the functions you just defined, they need to be added as layers using the [`tl.Fn()`](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.base.Fn) function.

# In[33]:

# UNQ_C5
# GRADED FUNCTION: CausalAttention
def CausalAttention(d_feature,
                    n_heads,
                    compute_attention_heads_closure=compute_attention_heads_closure,
                    dot_product_self_attention=dot_product_self_attention,
                    compute_attention_output_closure=compute_attention_output_closure,
                    mode='train'):
    """Transformer-style multi-headed causal attention.

    Args:
        d_feature (int):  dimensionality of feature embedding.
        n_heads (int): number of attention heads.
        compute_attention_heads_closure (function): Closure around compute_attention heads.
        dot_product_self_attention (function): dot_product_self_attention function.
        compute_attention_output_closure (function): Closure around compute_attention_output.
        mode (str): 'train' or 'eval'.

    Returns:
        trax.layers.combinators.Serial: Multi-headed self-attention model.
    """

    assert d_feature % n_heads == 0
    d_head = d_feature // n_heads

    ### START CODE HERE ###
    # (REPLACE INSTANCES OF 'None' WITH YOUR CODE)

    # HINT: The second argument to tl.Fn() is an uncalled function (without the parentheses)
    # Since you are dealing with closures you might need to call the outer
    # function with the correct parameters to get the actual uncalled function.
    ComputeAttentionHeads = tl.Fn('AttnHeads', compute_attention_heads_closure(n_heads, d_head), n_out=1)

    return tl.Serial(
        tl.Branch( # creates three towers for one input, takes activations and creates queries keys and values
            [tl.Dense(d_feature), ComputeAttentionHeads], # queries
            [tl.Dense(d_feature), ComputeAttentionHeads], # keys
            [tl.Dense(d_feature), ComputeAttentionHeads], # values
        ),

        tl.Fn('DotProductAttn', dot_product_self_attention, n_out=1), # takes QKV
        # HINT: The second argument to tl.Fn() is an uncalled function
        # Since you are dealing with closures you might need to call the outer
        # function with the correct parameters to get the actual uncalled function.
        tl.Fn('AttnOutput', compute_attention_output_closure(n_heads, d_head), n_out=1), # to allow for parallel
        tl.Dense(d_feature)
    )

    ### END CODE HERE ###

# In[34]:

# Take a look at the causal attention model
print(CausalAttention(d_feature=512, n_heads=8))

# **Expected Output:**
# ```CPP
# Serial[
#   Branch_out3[
#     [Dense_512, AttnHeads]
#     [Dense_512, AttnHeads]
#     [Dense_512, AttnHeads]
#   ]
#   DotProductAttn_in3
#   AttnOutput
#   Dense_512
# ]
# ```

# In[35]:

# UNIT TEST
# test CausalAttention
w2_tests.test_CausalAttention(CausalAttention)

#
#
# ## 2.3 Transformer decoder block
#
# Now that you have implemented the causal part of the transformer, you will implement the transformer decoder block. Concretely you will be implementing this image now.
#
#
#
# To implement this function, you will have to call the `CausalAttention` or Masked multi-head attention function you implemented above. You will have to add a feedforward which consists of:
#
# -  [tl.LayerNorm](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.normalization.LayerNorm) : used to layer normalize
# -  [tl.Dense](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Dense) : the dense layer
# -  [ff_activation](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.activation_fns.Relu) : feed forward activation (we use ReLu) here.
# -  [tl.Dropout](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Dropout) : dropout layer
# -  [tl.Dense](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Dense) : dense layer
# -  [tl.Dropout](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Dropout) : dropout layer
#
# Finally once you implement the feedforward, you can go ahead and implement the entire block using:
#
# -  [tl.Residual](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.combinators.Residual) : takes in the tl.LayerNorm(), causal attention block, tl.dropout.
#
# -  [tl.Residual](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.combinators.Residual) : takes in the feedforward block you will implement.
#
#
# ### Exercise 03
# **Instructions:** Implement the transformer decoder block. Good luck!

# In[36]:

# UNQ_C6
# GRADED FUNCTION: DecoderBlock
def DecoderBlock(d_model, d_ff, n_heads,
                 dropout, mode, ff_activation):
    """Returns a list of layers that implements a Transformer decoder block.

    The input is an activation tensor.

    Args:
        d_model (int):  depth of embedding.
        d_ff (int): depth of feed-forward layer.
        n_heads (int): number of attention heads.
        dropout (float): dropout rate (how much to drop out).
        mode (str): 'train' or 'eval'.
        ff_activation (function): the non-linearity in feed-forward layer.

    Returns:
        list: list of trax.layers.combinators.Serial that maps an activation tensor to an activation tensor.
    """

    ### START CODE HERE (REPLACE INSTANCES OF 'None' WITH YOUR CODE) ###

    # Create masked multi-head attention block using CausalAttention function
    causal_attention = CausalAttention(
                        d_model,
                        n_heads=n_heads,
                        mode=mode
                        )

    # Create feed-forward block (list) with two dense layers with dropout and input normalized
    feed_forward = [
        # Normalize layer inputs
        tl.LayerNorm(),
        # Add first feed forward (dense) layer (don't forget to set the correct value for n_units)
        tl.Dense(d_ff),
        # Add activation function passed in as a parameter (you need to call it!)
        ff_activation(), # Generally ReLU
        # Add dropout with rate and mode specified (i.e., don't use dropout during evaluation)
        tl.Dropout(rate=dropout, mode=mode),
        # Add second feed forward layer (don't forget to set the correct value for n_units)
        tl.Dense(d_model),
        # Add dropout with rate and mode specified (i.e., don't use dropout during evaluation)
        tl.Dropout(rate=dropout,mode=mode)
    ]

    # Add list of two Residual blocks: the attention with normalization and dropout and feed-forward blocks
    return [
      tl.Residual(
          # Normalize layer input
          tl.LayerNorm(),
          # Add causal attention block previously defined (without parentheses)
          causal_attention,
          # Add dropout with rate and mode specified
          tl.Dropout(rate=dropout, mode=mode)
        ),
      tl.Residual(
          # Add feed forward block (without parentheses)
          feed_forward
        ),
      ]
    ### END CODE HERE ###

# In[37]:

# Take a look at the decoder block
print(DecoderBlock(d_model=512, d_ff=2048, n_heads=8, dropout=0.1, mode='train', ff_activation=tl.Relu))

# **Expected Output:**
# ```CPP
# [Serial[
#   Branch_out2[
#     None
#     Serial[
#       LayerNorm
#       Serial[
#         Branch_out3[
#           [Dense_512, AttnHeads]
#           [Dense_512, AttnHeads]
#           [Dense_512, AttnHeads]
#         ]
#         DotProductAttn_in3
#         AttnOutput
#         Dense_512
#       ]
#       Dropout
#     ]
#   ]
#   Add_in2
# ], Serial[
#   Branch_out2[
#     None
#     Serial[
#       LayerNorm
#       Dense_2048
#       Serial[
#         Relu
#       ]
#       Dropout
#       Dense_512
#       Dropout
#     ]
#   ]
#   Add_in2
# ]]
# ```

# In[38]:

# UNIT TEST
# test DecoderBlock
w2_tests.test_DecoderBlock(DecoderBlock)

#
# ## 2.4 Transformer Language Model
#
# You will now bring it all together. In this part you will use all the subcomponents you previously built to make the final model. Concretely, here is the image you will be implementing.
#
#
#
#
# ### Exercise 04
# **Instructions:** Previously you coded the decoder block. Now you will code the transformer language model. Here is what you will need.
#
# -  positional_enconder - a list containing the following layers:
#     -  [tl.Embedding](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Embedding)
#     -  [tl.Dropout](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Dropout)
#     -  [tl.PositionalEncoding](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.attention.PositionalEncoding)
#
# - A list of `n_layers`  decoder blocks.
# -  [tl.Serial](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.combinators.Serial):  takes in the following layers or lists of layers:
#     -  [tl.ShiftRight](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.attention.ShiftRight): : shift the tensor to the right by padding on axis 1.
#     -  positional_encoder : encodes the text positions.
#     -  decoder_blocks : the ones you created.
#     -  [tl.LayerNorm](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.normalization.LayerNorm) : a layer norm.
#     -  [tl.Dense](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.Dense) : takes in the vocab_size.
#     -  [tl.LogSoftmax](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.core.LogSoftmax) : to predict.
#
# Go go go!! You can do it :)
#
#

# In[39]:

# UNQ_C7
# GRADED FUNCTION: TransformerLM
def TransformerLM(vocab_size=33300,
                  d_model=512,
                  d_ff=2048,
                  n_layers=6,
                  n_heads=8,
                  dropout=0.1,
                  max_len=4096,
                  mode='train',
                  ff_activation=tl.Relu):
    """Returns a Transformer language model.

    The input to the model is a tensor of tokens. (This model uses only the
    decoder part of the overall Transformer.)

    Args:
        vocab_size (int): vocab size.
        d_model (int):  depth of embedding.
        d_ff (int): depth of feed-forward layer.
        n_layers (int): number of decoder layers.
        n_heads (int): number of attention heads.
        dropout (float): dropout rate (how much to drop out).
        max_len (int): maximum symbol length for positional encoding.
        mode (str): 'train', 'eval' or 'predict', predict mode is for fast inference.
        ff_activation (function): the non-linearity in feed-forward layer.

    Returns:
        trax.layers.combinators.Serial: A Transformer language model as a layer that maps from a tensor of tokens
        to activations over a vocab set.
    """

    ### START CODE HERE (REPLACE INSTANCES OF 'None' WITH YOUR CODE) ###

    # Embedding inputs and positional encoder
    positional_encoder = [
        # Add embedding layer of dimension (vocab_size, d_model)
        tl.Embedding(vocab_size, d_model),
        # Use dropout with rate and mode specified
        tl.Dropout(rate=dropout, mode=mode),
        # Add positional encoding layer with maximum input length and mode specified
        tl.PositionalEncoding(max_len=max_len, mode=mode)]

    # Create stack (list) of decoder blocks with n_layers with necessary parameters
    decoder_blocks = [
        DecoderBlock(d_model, d_ff, n_heads,
                    dropout, mode, ff_activation) for _ in range(n_layers)]

    # Create the complete model as written in the figure
    return tl.Serial(
        # Use teacher forcing (feed output of previous step to current step)
        tl.ShiftRight(mode=mode), # Specify the mode!
        # Add positional encoder
        positional_encoder,
        # Add decoder blocks
        decoder_blocks,
        # Normalize layer
        tl.LayerNorm(),

        # Add dense layer of vocab_size (since need to select a word to translate to)
        # (a.k.a., logits layer. Note: activation already set by ff_activation)
        tl.Dense(vocab_size),
        # Get probabilities with Logsoftmax
        tl.LogSoftmax()
    )

    ### END CODE HERE ###

# In[40]:

# Take a look at the Transformer
print(TransformerLM(n_layers=1))

# **Expected Output:**
# ```CPP
# Serial[
#   Serial[
#     ShiftRight(1)
#   ]
#   Embedding_33300_512
#   Dropout
#   PositionalEncoding
#   Serial[
#     Branch_out2[
#       None
#       Serial[
#         LayerNorm
#         Serial[
#           Branch_out3[
#             [Dense_512, AttnHeads]
#             [Dense_512, AttnHeads]
#             [Dense_512, AttnHeads]
#           ]
#           DotProductAttn_in3
#           AttnOutput
#           Dense_512
#         ]
#         Dropout
#       ]
#     ]
#     Add_in2
#   ]
#   Serial[
#     Branch_out2[
#       None
#       Serial[
#         LayerNorm
#         Dense_2048
#         Serial[
#           Relu
#         ]
#         Dropout
#         Dense_512
#         Dropout
#       ]
#     ]
#     Add_in2
#   ]
#   LayerNorm
#   Dense_33300
#   LogSoftmax
# ]
# ```

# In[41]:

# UNIT TEST
# test TransformerLM
w2_tests.test_TransformerLM(TransformerLM)

#
# # Part 3: Training
#
# Now you are going to train your model. As usual, you have to define the cost function, the optimizer, and decide whether you will be training it on a `gpu` or `cpu`. In this case, you will train your model on a cpu for a few steps and we will load in a pre-trained model that you can use to predict with your own words.

#
# ### 3.1 Training the model
#
# You will now write a function that takes in your model and trains it. To train your model you have to decide how many times you want to iterate over the entire data set. Each iteration is defined as an `epoch`. For each epoch, you have to go over all the data, using your training iterator.
#
#
# ### Exercise 05
# **Instructions:** Implement the `train_model` program below to train the neural network above. Here is a list of things you should do:
#
# - Create the train task by calling [`trax.supervised.training.TrainTask`](https://trax-ml.readthedocs.io/en/latest/trax.supervised.html#trax.supervised.training.TrainTask) and pass in the following:
#     -  labeled_data  = train_gen
#     -  loss_fn  = [tl.CrossEntropyLoss()](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.metrics.CrossEntropyLoss)
#     -  optimizer  = [trax.optimizers.Adam(0.01)](https://trax-ml.readthedocs.io/en/latest/trax.optimizers.html#trax.optimizers.adam.Adam)
#     -  lr_schedule  = [lr_schedule](https://trax-ml.readthedocs.io/en/latest/trax.supervised.html#trax.supervised.lr_schedules.warmup_and_rsqrt_decay)
#
#
# - Create the eval task by calling [`trax.supervised.training.EvalTask`](https://trax-ml.readthedocs.io/en/latest/trax.supervised.html#trax.supervised.training.EvalTask) and pass in the following:
#     -  labeled_data  = eval_gen
#     -  metrics  = tl.CrossEntropyLoss() and [tl.Accuracy()](https://trax-ml.readthedocs.io/en/latest/trax.layers.html#trax.layers.metrics.Accuracy)
#
#
# - Create the training loop by calling [`trax.supervised.Training.Loop`](https://trax-ml.readthedocs.io/en/latest/trax.supervised.html#trax.supervised.training.Loop) and pass in the following:
#     -  TransformerLM
#     -  train_task
#     -  eval_task  = [eval_task]
#     -  output_dir = output_dir
#
# You will be using a cross entropy loss, with Adam optimizer. Please read the [Trax](https://trax-ml.readthedocs.io/en/latest/index.html) documentation to get a full understanding.
#
# The training loop that this function returns can be runned using the `run()` method by passing in the desired number of steps.

# In[42]:

from trax.supervised import training

# UNQ_C8
# GRADED FUNCTION: train_model
def training_loop(TransformerLM, train_gen, eval_gen, output_dir = "~/model"):
    '''
    Input:
        TransformerLM (trax.layers.combinators.Serial): The model you are building.
        train_gen (generator): Training stream of data.
        eval_gen (generator): Evaluation stream of data.
        output_dir (str): folder to save your file.

    Returns:
        trax.supervised.training.Loop: Training loop.
    '''
    output_dir = os.path.expanduser(output_dir)  # trainer is an object
    lr_schedule = trax.lr.warmup_and_rsqrt_decay(n_warmup_steps=1000, max_value=0.01)

    ### START CODE HERE (REPLACE INSTANCES OF 'None' WITH YOUR CODE) ###
    train_task = training.TrainTask(
      labeled_data=train_gen, # The training generator
      loss_layer=tl.CrossEntropyLoss(), # Loss function (Don't forget to instantiate!)
      optimizer=trax.optimizers.Adam(0.01), # Optimizer (Don't forget to set LR to 0.01)
      lr_schedule=lr_schedule,
      n_steps_per_checkpoint=10
    )

    eval_task = training.EvalTask(
      labeled_data=eval_gen, # The evaluation generator
      metrics=[tl.CrossEntropyLoss(), tl.Accuracy()] # CrossEntropyLoss and Accuracy (Don't forget to instantiate both!)
    )

    ### END CODE HERE ###

    loop = training.Loop(TransformerLM(d_model=4,
                                       d_ff=16,
                                       n_layers=1,
                                       n_heads=2,
                                       mode='train'),
                         train_task,
                         eval_tasks=[eval_task],
                         output_dir=output_dir)

    return loop

# Notice that the model will be trained for only 10 steps.
#
# Even with this constraint the model with the original default arguments took a very long time to finish. Because of this some parameters are changed when defining the model that is fed into the training loop in the function above.

# In[43]:

# UNIT TEST
# test training_loop
w2_tests.test_training_loop(training_loop, TransformerLM)

# In[44]:

# Should take around 1.5 minutes
get_ipython().system('rm -f ~/model/model.pkl.gz')
loop = training_loop(TransformerLM, train_batch_stream, eval_batch_stream)
loop.run(10)

#
#  # Part 4:  Evaluation
#
#
# ### 4.1 Loading in a trained model
#
# In this part you will evaluate by loading in an almost exact version of the model you coded, but we trained it for you to save you time. Please run the cell below to load in the model.
#
# As you may have already noticed the model that you trained and the pretrained model share the same overall architecture but they have different values for some of the parameters:
#
#
#    `Original (pretrained) model: `
#
#     TransformerLM(vocab_size=33300, d_model=512, d_ff=2048, n_layers=6, n_heads=8,
#                    dropout=0.1, max_len=4096, ff_activation=tl.Relu)
#
#    `Your model:`
#
#     TransformerLM(d_model=4, d_ff=16, n_layers=1, n_heads=2)
#
#    **Only the parameters shown for your model were changed. The others stayed the same.**

# In[45]:

# Get the model architecture
model = TransformerLM(mode='eval')

# Load the pre-trained weights
model.init_from_file('model.pkl.gz', weights_only=True)

#
# # Part 5: Testing with your own input
#
# You will now test your input. You are going to implement greedy decoding. This consists of two functions. The first one allows you to identify the next symbol. It gets the argmax of the output of your model and then returns that index.
#
#
# ### Exercise 06
# **Instructions:** Implement the next symbol function that takes in the cur_output_tokens and the trained model to return the the index of the next word.

# In[48]:

# UNQ_C9
def next_symbol(cur_output_tokens, model):
    """Returns the next symbol for a given sentence.

    Args:
        cur_output_tokens (list): tokenized sentence with EOS and PAD tokens at the end.
        model (trax.layers.combinators.Serial): The transformer model.

    Returns:
        int: tokenized symbol.
    """
    ### START CODE HERE (REPLACE INSTANCES OF 'None' WITH YOUR CODE) ###

    # current output tokens length
    token_length = len(cur_output_tokens)
    # calculate the minimum power of 2 big enough to store token_length
    # HINT: use np.ceil() and np.log2()
    # add 1 to token_length so np.log2() doesn't receive 0 when token_length is 0
    padded_length = 2**int(np.ceil(np.log2(token_length + 1)))

    # Fill cur_output_tokens with 0's until it reaches padded_length
    padded = cur_output_tokens + [0] * (padded_length - token_length)
    padded_with_batch = np.array(padded)[None, :] # Don't replace this None! This is a way of setting the batch dim

    # model expects a tuple containing two padded tensors (with batch)
    output, _ = model((padded_with_batch, padded_with_batch))
    # HINT: output has shape (1, padded_length, vocab_size)
    # To get log_probs you need to index output wih 0 in the first dim
    # token_length in the second dim and all of the entries for the last dim.
    log_probs = output[0, token_length, :]

    ### END CODE HERE ###

    return int(np.argmax(log_probs))

# In[49]:

# Test it out!
sentence_test_nxt_symbl = "I want to fly in the sky."
detokenize([next_symbol(tokenize(sentence_test_nxt_symbl)+[0], model)])

# **Expected Output:**
# ```CPP
# 'The'
# ```

# In[50]:

# UNIT TEST
# test next_symbol
w2_tests.test_next_symbol(next_symbol, TransformerLM)

#
# ### 5.1 Greedy decoding
#
# Now you will implement the greedy_decode algorithm that will call the `next_symbol` function. It takes in the input_sentence, the trained model and returns the the decoded sentence.
#
#
# ### Exercise 07
#
# **Instructions**: Implement the greedy_decode algorithm.

# In[51]:

# UNQ_C10
# Decoding functions.
def greedy_decode(input_sentence, model, next_symbol=next_symbol, tokenize=tokenize, detokenize=detokenize):
    """Greedy decode function.

    Args:
        input_sentence (string): a sentence or article.
        model (trax.layers.combinators.Serial): Transformer model.

    Returns:
        string: summary of the input.
    """

    ### START CODE HERE (REPLACE INSTANCES OF 'None' WITH YOUR CODE) ###
    # Use tokenize()
    cur_output_tokens = tokenize(input_sentence) + [0]
    generated_output = []
    cur_output = 0
    EOS = 1

    while cur_output != EOS:
        # Get next symbol
        cur_output = next_symbol(cur_output_tokens, model)
        # Append next symbol to original sentence
        cur_output_tokens.append(cur_output)
        # Append next symbol to generated sentence
        generated_output.append(cur_output)
        print(detokenize(generated_output))

    ### END CODE HERE ###

    return detokenize(generated_output)

# In[52]:

# Test it out on a sentence!
test_sentence = "It was a sunny day when I went to the market to buy some flowers. But I only found roses, not tulips."
print(wrapper.fill(test_sentence), '\n')
print(greedy_decode(test_sentence, model))

# **Expected Output:**
# ```CPP
# :
# : I
# : I just
# : I just found
# : I just found ros
# : I just found roses
# : I just found roses,
# : I just found roses, not
# : I just found roses, not tu
# : I just found roses, not tulips
# : I just found roses, not tulips
# : I just found roses, not tulips.
# : I just found roses, not tulips.<EOS>
# : I just found roses, not tulips.<EOS>
# ```

# In[53]:

# Test it out with a whole article!
article = "It’s the posing craze sweeping the U.S. after being brought to fame by skier Lindsey Vonn, soccer star Omar Cummings, baseball player Albert Pujols - and even Republican politician Rick Perry. But now four students at Riverhead High School on Long Island, New York, have been suspended for dropping to a knee and taking up a prayer pose to mimic Denver Broncos quarterback Tim Tebow. Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were all suspended for one day because the ‘Tebowing’ craze was blocking the hallway and presenting a safety hazard to students. Scroll down for video. Banned: Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll (all pictured left) were all suspended for one day by Riverhead High School on Long Island, New York, for their tribute to Broncos quarterback Tim Tebow. Issue: Four of the pupils were suspended for one day because they allegedly did not heed to warnings that the 'Tebowing' craze at the school was blocking the hallway and presenting a safety hazard to students."
print(wrapper.fill(article), '\n')
print(greedy_decode(article, model))

# **Expected Output:**
# ```CPP
# Jordan
# Jordan Ful
# Jordan Fulcol
# Jordan Fulcoly
# Jordan Fulcoly,
# Jordan Fulcoly, Wayne
# Jordan Fulcoly, Wayne Dre
# Jordan Fulcoly, Wayne Drexe
# Jordan Fulcoly, Wayne Drexel
# Jordan Fulcoly, Wayne Drexel,
# .
# .
# .
#
# Final summary:
#
# Jordan Fulcoly, Wayne Drexel, Tyler Carroll and Connor Carroll were
# suspended for one day. Four students were suspended for one day
# because they allegedly did not heed to warnings that the 'Tebowing'
# craze was blocking the hallway and presenting a safety hazard to
# students.<EOS>
# ```

# In[54]:

# UNIT TEST
# test greedy_decode
w2_tests.test_greedy_decode(greedy_decode)

# **Congratulations on finishing this week's assignment!** You did a lot of work and now you should have a better understanding of the enconder part of Transformers and how Transformers can be used for text summarization.
#
# **Keep it up!**


#!/usr/bin/env python3

import unittest
from unittest.mock import patch
import pandas as pd
import numpy as np

from tmc import points

from tmc.utils import load, get_out, patch_helper

module_name="src.operations_on_series"
create_series = load(module_name, "create_series")
modify_series = load(module_name, "modify_series")
main = load(module_name, "main")
ph = patch_helper(module_name)

@points('p03-14.1')
class OperationsOnSeries(unittest.TestCase):

    def test_creation(self):
        self.assertEqual("", "")
        L1=[2,3,4]
        L2=[9,8,7]
        indices=list("abc")
#        with patch(patch_name(module_name, "pd.core.series.Series"), wraps=pd.core.series.Series) as ps:
        with patch(ph("pd.Series"), wraps=pd.Series) as ps:
            ret = create_series(L1, L2)
            self.assertEqual(len(ret), 2, msg="Expected a pair of Series as a return value from function create_series!")
            s1, s2 = ret
            #ps.assert_called()
            self.assertEqual(ps.call_count, 2, msg="Expected the constructor pd.Series to be called exactly twice!")
        np.testing.assert_array_equal(s1.values, L1,
                                      err_msg="Expected values of first series to be %s" % L1)
        np.testing.assert_array_equal(s2.values, L2,
                                      err_msg="Expected values of second series to be %s" % L2)
        np.testing.assert_array_equal(s1.index, indices,
                                      err_msg="Expected the index of first series to be %s" % indices)
        np.testing.assert_array_equal(s2.index, indices,
                                      err_msg="Expected the index of second series to be %s" % indices)

    def test_modification(self):
        indices=list("abc")
        s1 = pd.Series([0,1,2], index=indices)
        s2 = pd.Series([3,4,5], index=indices)
        ret = modify_series(s1, s2)
        self.assertEqual(len(ret), 2, msg="Expected modify_series to return a pair of Series!")
        t1, t2 = ret
        self.assertIsInstance(t1, pd.Series, msg="Expected modify_series to return a pair of Series!")
        self.assertIsInstance(t2, pd.Series, msg="Expected modify_series to return a pair of Series!")
        t1_ind=list("abcd")
        t2_ind=list("ac")
        np.testing.assert_array_equal(t1.index, t1_ind,
                                      err_msg="Expected the index of first series to be %s!" % t1_ind)
        np.testing.assert_array_equal(t2.index, t2_ind,
                                      err_msg="Expected the index of second series to be %s!" % t2_ind)

        np.testing.assert_array_equal(t1.values, [0,1,2,4],
                                      err_msg="Values of first series is not correct!")
        np.testing.assert_array_equal(t2.values, [3, 5],
                                      err_msg="Values of second series is not correct!")

    def test_main(self):
        with patch(ph("create_series"), wraps=create_series) as pcs,\
             patch(ph("pd.Series.__add__"), side_effect=[pd.Series()]) as padd,\
             patch(ph("modify_series"), wraps=modify_series) as pms:
            main()
            pcs.assert_called()
            pms.assert_called()
            padd.assert_called()

if __name__ == '__main__':
    unittest.main()


#!/usr/bin/env python
# coding: utf-8

# # Lab 1 Solutions
#
# ### Objectives
# In this lab, we'll
# - Review the computational infrastructure around our data science environments,
# - Go through the process of ensuring that we have a Python environment set up for this class with the proper installed packages
# - Within our environment, we'll review the basic data science operations in Python, and introduce some tips and tricks.
#
#
# ```{admonition} Take a Deep Breath
# Don't freak out if stuff presented here is brand new to you! Ask a friend, google around (esp. stack overflow) and find a solution. All of these examples can be done in a few lines of code.
# ```
#
# # Part I: Computational Ecosystem
#
# In the space below, or in your own assignment, answer the following:
#
# ## Question A
#
# Describe the following terms, and point out the differences between them. Feel free to look things up.
# - python:
# - the terminal:
# - the file system:
# - jupyter:
# - an IDE:
# - a text editor:
# - git:
# - PATH:
#

# I am writing this lab in a notebook. We'll be discussing the pros and cons of notebooks in class. Below, I'm going to check which python installation on my computer is being pointed to within my PATH. As a reminder, you can see your full path by echoing it from the terminal. Within a notebook, that looks like this:

# In[2]:

get_ipython().system('echo $PATH')

# ```{note}
# The "!" in my notebook allows me to run terminal commands from a notebook; you don't need this symbol when running commands in an actual terminal.
# ```

# I can check my python as follows:

# In[3]:

get_ipython().system('which python')

# We can see that calls to `python` are triggering the python installed in `anaconda3`, which is what we want (see the installation video for more details). If your call to `which python` in the terminal returns something like `usr/bin/python`, then something has likely gone wrong with your installation. There are some troubleshooting steps suggested in the installation video.
#
# ## Setting up an Environment.
#
# We can think of packages as programs installed on our computer. But what if one of my projects needs Photoshop 14.0, and another needs a feature that was only available in Photoshop 12.5.2? An environment is the system on which your code is being executed. When you fire up a terminal, this is usually your *base* environment, the default one for your *user* on a given computer system. But rather than always installing programs in this base installation, we can create custom environments for each of our projects. We can then install the exact dependencies for those projects within our environments, and we'll know they won't mess with each other.
#
# For this class, we're going to be doing a lot of package installations. To ensure we are all on the same page and are working with the same tools, we're going to use a `conda environment` to maintain versioning. Note: there are multiple environment creation tools/methods. Conda environments are the predominant standard in astronomy, hence their use here.
#
# In your terminal, type the following:
# ```{note}
# If you are on WINDOWS, you NEED to use the ANACONDA TERMINAL. NOT YOUR WINDOWS POWER SHELL/COMMAND PROMPT. Search your pc for anaconda and you'll see an anaconda launcher (if you followed the installation video correctly). From there, you should be able to find an anaconda terminal/prompt. That's where you should do anything whenever I say "from your terminal".
# ```
#
#

# In[ ]:

conda create -n a330 python=3.8

# Once you run this, answer "y" to the prompts, and your new environment will be installed.
#
# ```{note}
# The above command may take several minutes to execute.
# ```
#
#
# Next, we want to activate this environment (still in our terminal). We do this as follows:

# In[ ]:

conda activate a330

# When you do so, you should see the left hand edge of your prompt switch from (base) to (a330).
#
# Next, let's make an alias so that getting into our a330 environment is a snap. We're going to access a file called `.bash_profile`, which allows us to set aliases and environment variables. This file is located in your home directory, so I can print mine here:

# In[6]:

get_ipython().system('more ~/.bash_profile')

# Notice above I use the `~` which is a shorthand for home directory. On my computer, the default home directory for my user is `/Users/ipasha/`.
#
#
#
# This file has some conda stuff in it at the top, as well as some path and python path exports, as well as an alias.
#
# Yours should also have the conda init stuff, if you installed anaconda properly. Using your text editor of choice, add a line to this file that reads `alias a330='conda activate a330'`.
#
# ```{sidebar} Using Vi/vim
# Vi/vim is a built-in terminal program that allows for the editing of files. It is helpful to learn, especially when working on remote servers. We'll go into it more later, but here is a step by step for performing the above step with vim.
# - First: from the terminal, type `vim ~/.bash_profile` and hit enter. This will open the editor. If 'vim' isn't recognized, try 'vi'.
# - Next: Press the "I" key to open insert mode. Move your cursor with the arrow keys to the desired line, then type in the alias command shown to left.
# - Finally: Press `esc` to get out of insert mode, then type `:wq` and hit enter in order to "write" then "quit".
# ```
#
#
#
# Now, from your terminal, source your profile by typing `source ~/.bash_profile`. You're good to go! Test that you can activate your environment by typing `a330` and hitting enter.
#
# ```{note}
# To deactivate, just type `conda deactivate`.
# ```
#
#

# ## Adding Jupyter
#
# You'll be using notebooks during this class, and we need to make sure that we can access our new environment from within Jupyter notebook. To ensure this, we're going to do the following:
#
# First, make sure your environment is activated.
#
# Then, type:

# In[ ]:

conda install -c anaconda ipykernel

# This ensures we can select different kernels inside jupyter. A kernel is basically "the thing that is python", the root thing being run on your system when you use python. By creating environments, we're creating different unique kernels, and we can now get to them within our notebooks.
#
# Now, run the following:

# In[ ]:

python -m ipykernel install --user --name=a330

# Once you've done this, you should have the ability to access your new environment from within Jupyter. We can test this as follows:
# - First, open a new terminal window, and activate your environment (if you made the alias, this means typing `a330` in your terminal.
# - Next, type `jupyter lab` to open jupyter lab. If for some reason you don't have jupyter lab yet, you can install it now with `conda install -c conda-forge jupyterlab`.
# - Once you have lab open, there should be a 'launcher' page, with one option being to create a new notebook using python -- you *should* see your environment listed there.
# - If you don't hit refresh on the webpage just in case.
# - You can also click on the option to open a python3 notebook. Inside, in the top right corner, it should say your current environment (probably Python 3). Clicking that, it should give you the option to choose a different environment, and your environment should be listed there.
#
# ```{note}
# If you already had a lab open, you'll have to hit refresh to get it to show up.
# ```
# ## Installing Packages
#
# Now that we have our environment, we're going to install the set of packages we need for this class. We may need more of them as the semester goes on, but for now, do the following (in your terminal, within your environment).

# In[ ]:

conda install -n a330 numpy scipy astropy matplotlib

# (again, hitting "y" when prompted). Again, this step might take a minute or so to run.
#
# Congrats, you now have an environment set up for this class, and can jump in and out of it at will, either in your terminal, or within a Jupyter notebook.
#
# ```{admonition} Hot Tip
# It's highly recommended you do these steps anytime you start a new research project. Up front, you may not know all the dependencies that will arise, but as you go along, if you keep your work to that environment, you'll be able to carefully control which versions of which packages you're accessing at all times.
# ```
#
#

# # Part II: Python Review
#
# In this section, I'll ask you to perform some pythonic operations to get back into the swing of things if it has been a little while.
#
# For this assignment, please carry out your work in a Jupyter notebook, with the questions labeled and your output shown. You'll submit this notebook via Github, but we will discuss how to perform this step in class.

# ## Question 1
# Create a 2D array of dimensions 1000 x 1000, in which the values in each pixel are random-gaussian distributed about a mean of 10, with a sigma of 2, and then use matplotlib to display this image. Make sure (0,0) is in the lower lefthand corner.

# In[145]:

import numpy as np
import scipy.stats as stats
import matplotlib.pyplot as plt

# In[146]:

# Solution 1
mu = 10
sigma = 2
Gauss = stats.norm(loc=mu,scale=sigma)
array = Gauss.rvs((1000,1000))
fig, ax = plt.subplots(figsize=(10,10))
ax.imshow(array,origin='lower');

# In[147]:

#Solution 2:
array_np = np.random.normal(loc=10,scale=2,size=(1000,1000))
fig, ax = plt.subplots(figsize=(10,10))
ax.imshow(array_np,origin='lower');

# ## Question 2
#
# The distribution of pixels in your above image should not have many outliers beyond 3-sigma from the mean, but there will be some. Find the location of any 2 sigma outliers in the image, and highlight them by circling their location.
# Confirm that the fraction of these out of the total number of pixels agrees with the expectation for a normal distribution.

# In[149]:

# Solution
outliers = np.where((array>mu+3*sigma)|(array<mu-3*sigma))
fig, ax = plt.subplots(figsize=(10,10))
ax.imshow(array,origin='lower')
ax.plot(outliers[0],outliers[1],'o',ms=5,color='None',mec='r');

# ## Question 3
#
# When dealing with astronomical data, it is sometimes advisable to not include outliers in a calculation being performed on a set of data (in this example, an image). We know, of course, that the data we're plotting ARE coming from a gaussian distribution, so there's no reason to exclude, e.g., 3-sigma outliers, but for this example, let's assume we want to.
#
# Create a numpy masked array in which all pixels that are > 3$\sigma$ from the image mean are masked. Then, calculate the mean and sigma of the new masked array.

# In[34]:

# Solution
clipped_array = np.ma.masked_where((array>mu+3*sigma)|(array<mu-3*sigma),array)
print(f'Clipped Mean: {np.mean(clipped_array):.3f}  |  Clipped Sigma: {np.std(clipped_array):.3f}')

# In[154]:

# solution 2
m= (array < mu+3.*sigma) & (array>mu-3.*sigma)

mn = np.mean(array[m])
st = np.std(array[m])
print('Clipped mean; {:0.3f};  Clipped Sigma: {:0.3f}'.format(mn,st))

# As expected, clipping the outliers of this distribution does not affect the mean in any strong way, but does noticably decrease $\sigma$.

# ## Question 4:
#
# Using Array indexing, re-plot the same array from above, but zoom in on the inner 20% of the image, such that the full width is 20% of the total. Note: try not to hard code your indexing. You should be able to flexibly change the percentage. For this one, use a white-to-black color map.
#

# In[150]:

# solution
cent = int(array.shape[0] / 2)
perc = int(0.2 * array.shape[0]*0.5)
cropped_array = array[cent-perc:cent+perc,cent-perc:cent+perc]
fig, ax = plt.subplots(figsize=(10,10))
ax.imshow(cropped_array,origin='lower',cmap='gray_r');

# As expected, our image is now 200 by 200 pixels across. Note that our new image has its own indexing. A common "gotcha" when working with arrays like this is to index in, but then try to use indices found (e.g., via `where()`) in the larger array on the cropped in version, which can lead to errors.

# ## Question 5
#
# Often, we have an expression to calculate of the form
#
# $$
# \sum_i \sum_j a_i b_j
# $$

# Your natural impulse for coding this double sum might look like this:

# In[ ]:

total = 0
for i in a:
    for j in b:
         total+= i*j

# which, mathematically, makes sense! But as it turns out, there's a way we can do this without any loops at all --- and when $\vec{a}$ and $\vec{b}$ get long, this becomes hugely important in our code.
#
# The trick we're going to use here is called [array broadcasting](https://numpy.org/doc/stable/user/basics.broadcasting.html), which you can read about at the link if you're not already familar. I'm going to give you $\vec{a}$ and $\vec{b}$ below. For this exercise, calculate the double sum indicated above without the use of a for-loop.
#
# ```{admonition} Hint
# The command `np.newaxis` will be useful here, or for a slightly longer solution, try `np.repeat` and `reshape()`.
# ```

# In[35]:

a = np.array([1,5,10,20])
b = np.array([1,2,4,16])

# Solution
output = np.sum(a[:,np.newaxis]*b)

# In[37]:

output

# We can confirmed the above worked using our slow loop:

# In[40]:

total = 0
for i in a:
    for j in b:
         total+= i*j
total

# We can also perform this trick without knowing about `np.newaxis` explicitly. You may have done it this way.

# In[55]:

matrix_a = np.repeat(a,len(b)).reshape(len(a),len(b))
matrix_b = np.repeat(b,len(a)).reshape((len(b)),len(a)).T
output = np.sum(matrix_a*matrix_b)
output

# I want to take a moment here and really highlight the difference in speed between these two methods. Let's bump up the length... by a lot:

# In[68]:

a = np.random.random(5000)
b = np.random.random(5000)

# In[69]:

get_ipython().run_cell_magic('timeit', '', '\ntotal = 0\nfor i in a:\n    for j in b:\n         total+= i*j')

# We can see that above, it took 5 whole seconds to run the cell. That's for only 5000 entries in each list. Many astronomical lists on which we might have to do this are 10,000 or more long! Notice this is an exponential operation. Try running the above on your computer with 10000 long lists. I'll wait... until you get bored and kill the process after many minutes.
#
# Now see below:

# In[70]:

get_ipython().run_cell_magic('timeit', '', 'output = np.sum(a[:,np.newaxis]*b)')

# Wow. 52 ms to do that whole thing. Broadcasting. We love it, we use it, we love it.
#
# ```{admonition} Takeaway Point
# The takeaway from this lesson should be that anytime we can do something without a loop... WE SHOULD.
# ```
#
#

# ## Question 6
#
# Often in astronomy we need to work with grids of values. For example, let's say we have a model that describes some data, and the model has 2 parameters, $a$ and $b$.
#
# We might choose different combinations of $a$ and $b$, and determine a metric for how well models of such combinations fit our data (e.g., $\chi^2$).
#
# We may then want to plot this $\chi^2$ value for each point on our grid -- that is, at each grid position corresponding to some $a_i$ and $b_j$.
#
# Below, I provide a function, `chi2`, which returns a single number given some singular inputs `a` and `b`.
#
# Create some arrays of `a` and `b` to test that range between 1 and 25, and have 10 entries evenly spaced between those values. Then, loop over them and find the $\chi^2$ using my function.
# ```{note}
# We can't get around the double loop in this case, because we are operating under the assumption that the calculation of some single $\chi^2$ using a unique combination of $a_i$ and $b_j$ cannot be vectorized. If it could, we wouldn't need to do this activity. But often, we can't, because the creation of a model given some inputs is nontrivial.
# ```
#
# Once you've stored the $\chi^2$ values for each combination of $a$ and $b$, create a plot with $a$ and $b$ as the axes and show using colored circles the $\chi^2$ value at each location. Add a colorbar to see the values being plotted.
#
# To create this grid, use the `np.meshgrid()` function. For your plot, make sure the marker size is big enough to see the colors well.
#
#
#
#

# In[151]:

def chi2(a,b):
    return ((15-a)**2+(12-b)**2)**0.2 #note, this is nonsense, but should return a different value for each input a,b

# Solution
a = np.linspace(1,25,10)
b = np.linspace(1,25,10)
chi2_values = []
for i in a:
    for j in b:
        chi2_values.append(chi2(i,j))

chi2_values = np.array(chi2_values)

xx,yy = np.meshgrid(a,b)

fig, ax = plt.subplots(figsize=(12,10))

im = ax.scatter(xx,yy,c=chi2_values,marker='o',s=100)
plt.colorbar(im);

# ## Question 7
#
# Re-show your final plot above, making the following changes:
#
# - label your colorbar as $\chi^2$ using latex notation, with a fontsize>13
# - Make your ticks point inward and be longer
# - Make your ticks appear on the top and right hand axes of the plot as well
# - If you didn't already, label the x and y axes appropriately and with a font size > 13
# - Make sure the numbers along the axes have fontsizes > 13
#

# In[152]:

fig, ax = plt.subplots(figsize=(12,10))

im = ax.scatter(xx,yy,c=chi2_values,marker='o',s=100)
ax.tick_params(direction='in',right=True,top=True,length=7,width=1.5,labelsize=14)
ax.set_xlabel(r'parameter $a$',fontsize=15)
ax.set_ylabel(r'parameter $b$',fontsize=15)
cbar = plt.colorbar(im)
cbar.set_label(r'model $\chi^2$',fontsize=14);

# ## Question 8
#
# Some quick list comprehensions! For any unfamilar, **comprehensions** are pythonic statements that allow you to compress a for-loop (generally) into a single line, and usually runs faster than a full loop (but not by a ton).
#
# Take the for-loop below and write it as a list comprehension.

# In[119]:

visited_cities = ['San Diego', 'Boston', 'New York City','Atlanta']
all_cities = ['San Diego', 'Denver', 'Boston', 'Portland', 'New York City', 'San Francisco', 'Atlanta']

not_visited = []
for city in all_cities:
    if city not in visited_cities:
        not_visited.append(city)

print(not_visited)

# In[121]:

# Solution

not_visited = [i for i in all_cities if i not in visited_cities]
print(not_visited)

# Next, create an array of integers including 1 through 30, inclusive. Using a comprehension, create a numpy array containing the squared value of only the odd numbers in your original array. (*Hint, remember the modulo operator*)

# In[122]:

# Solution
full = np.arange(1,31)
squared_odds = np.array([i**2 for i in full if i%2!=0])
print(squared_odds)

# In the next example, you have a list of first names and a list of last names. Use a list comprehension to create an array that is a list of full names (with a space between first and last names).

# In[123]:

first_names = ['Bob','Samantha','John','Renee']
last_names = ['Smith','Bee','Oliver','Carpenter']

# Solution
full_names = [i+' '+j for i,j in zip(first_names,last_names)]
print(full_names)

# ```{admonition} Challenge Problem (worth Extra Credit)
# I've created new lists that contain strings of the names in the format Lastname,Firstname, with random leading/trailing spaces and terrible capitalizations. Use a list comprehension to make our nice, "Firstname Lastname" list again.
# ```

# In[129]:

all_names = ['sMitH,BoB   ', '  bee,samanthA',' oLIVER,JOHN ','  caRPENTer,reneE  ']

# solution
full_names = [i.strip().split(',')[1].upper()[0]
              +i.strip().split(',')[1].lower()[1:]
              +' '
              +i.strip().split(',')[0].upper()[0]
              +i.strip().split(',')[0].lower()[1:]
              for i in all_names]
print(full_names)

# ```{note}
# Note that with this last example, we're entering a degree of single-line length and complexity that it almost doesn't make sense to use a comprehension anymore. Just because something CAN be done in one line doesn't mean is has to be, or should be.
#
# ```
#
# You may be wondering what use case this type of coding has in astronomy -- turns out, quite a lot. Take this example: you read in a data table and the columns have names like "FLUX HA", "FLUX ERR", etc.
#
# If you're trying to make a `pandas DataFrame` of this table, it is advantageous to rename these columns something like `flux_ha` and `flux_err`. This way, commands like `df.flux_ha` can be used.
#
# Being able to iterate over the string list of column names and turn caps into lower case, spaces into underscores, etc., is a useful skill that will come in handy when wrangling data.
#
# Below, for reference, I show how I myself would do the above example in production code:

# In[133]:

def clean_csv_string(str_in,sep=',',formatting='LastFirst'):
    str_stripped = str_in.strip() # remove trailing/leading spaces
    str_split = str_stripped.split(sep) #split at delimiter
    str_cap_correct = [i.upper()[0]+i.lower()[1:] for i in str_split]
    if formatting=='LastFirst':
        out_string = str_cap_correct[1] + ' ' + str_cap_correct[0]
        return out_string
    elif formmatting=='FirstLast':
        out_string = str_cap_correct[0] + ' ' + str_cap_correct[1]
        return out_string
    else:
        raise ValueError('formatting not set correctly')

full_names = [clean_csv_string(i) for i in all_names]
print(full_names)

# By making the string cleaning steps a function, I could take the time to explain what is going on within the function, as well as control for additional possibilities (like the names being in First,Last formatting). I could make this function more robust and complex, and my final comprehension stays readable and simple, as I loop over the names and run each through my handy functions (with some settings tweaked, potentially).

# ## Question 9
#
# Take the arrays `XX`, `YY`, and `ZZ` below and create one multidimensional array in which they are the columns. Print to confirm this worked.

# In[131]:

XX = np.array([1,2,3,4,5,6,7,8,9])
YY = np.array([5,6,7,8,9,10,11,12,13])
ZZ = np.array([10,11,12,13,14,15,16,17,18])

#solution
cols = np.column_stack((XX,YY,ZZ))
cols

# ## Question 10
#
# Units, units, units. The bane of every scientists' existence... except theorists that set every constant equal to 1.
#
# In the real world, we measure fluxes or magnitudes in astronomical images, infer temperatures and densities from data and simulations, and ultimately have to deal with units one way or another.
#
# Thankfully, our friends at `astropy` know this, and they've come to save the day. This next question serves as an introduction to the `units` module in astropy, which can be both a live saver and a pain in the ass, but at the end of the day is absolutely worth learning.

# In[137]:

import astropy.units as u

# The standard import for this library is `u`, so be careful not to name any variables that letter.
#
# To "assign" units to a variable, we multiply by the desired unit as follows. Note that generally the module knows several aliases/common abrreviations for a unit, if it is uniquely identifiable.

# In[138]:

star_temp = 5000*u.K
star_radius = 0.89 * u.Rsun
star_mass = 0.6 * u.Msun

# We can perform trivial conversions using the `.to()` method.

# In[139]:

star_radius.to(u.km)

# Once we attach units to something, it is now a `Quantity` object. Quantity objects are great, above, we saw they have built-in methods to facilitate conversion. They can also be annoying -- sometimes another function we've written needs just the raw value or array back out. To get this, we use the `.value` attribute of a quantity object:

# In[140]:

star_mass.to(u.kg).value

# This now strips away all `Quantity` stuff and gives us an array or value to use elsewhere in our code.
#
# Units are great because they help us combine quantities while tracking units and dimensional analysis. A common operation in astronomy is converting a flux to a luminosity given a distance, using
#
# $$
# F = \frac{L}{4\pi D^2}
# $$
# where $L$ is the luminosity and $D$ is the distance to the source.
#
# What if I've made a flux measurement in astronomical units such as erg/s/cm$^2$, and I want to know the luminosity in solar luminosities, and my distance happens to be in Mpc? Regardless of my input units, I can easily do this:

# In[141]:

L = 4 * np.pi * (3.6*u.Mpc)**2 * (7.5e-14 * u.erg/u.s/u.cm**2)
L.to(u.Lsun)

# This conversion worked because the units worked out. If my units of flux weren't correct, I'd get an error:

# In[142]:

L = 4 * np.pi * (3.6*u.Mpc)**2 * (7.5e-14 * u.erg/u.s/u.cm**2/u.AA)
L.to(u.Lsun)

# Here, `units` realized that I was putting in units of flux density, but wanted a luminosity out, and ultimately those units don't resolve out. Thus, it can be a great way to catch errors in your inputs.
#
# Note: just be careful that sometimes, you throw a constant into an equation but the constant has some units. If you're going to use the unit module to do a calculation, ALL inputs that HAVE units must be assigned them correctly as above for it to work.

# For your exercise, consider the following:
#
# The virial temperature of a galaxy halo is given roughly by
#
# $$
# T_{\rm vir} \simeq 5.6\times10^4\;\textrm{K}\left(\frac{\mu}{0.59}\right)\left(\frac{M_{\rm halo}}{10^{10}\; M_{\odot}}\right)^{2/3}\left(\frac{1+z}{4}\right)
# $$

# where here, we can assume $\mu$ is 0.59.
#
# Write a function that takes as an input a halo mass, redshift, and optionally $\mu$ (default 0.59), and returns the virial temperature in Kelvin. Your function should take in an astropy quantity with mass units, but should allow for the mass to be input with any appropriate units.

# In[144]:

# Solution

def Tvir(M,z,mu=0.59):
    return ((5.6e4*u.K)*(mu/0.59)*(M/(1e10*u.Msun))**(2/3.)*((1+z)/4)).to(u.K)

Tvir(10**11*u.Msun,z=3)

import numpy,import scipy,import astropy


"""
Defines the unit tests for the
:mod:`colour.models.rgb.transfer_functions.itur_bt_601` module.
"""

import numpy as np
import unittest

from colour.models.rgb.transfer_functions import oetf_BT601, oetf_inverse_BT601
from colour.utilities import domain_range_scale, ignore_numpy_errors

__author__ = "Colour Developers"
__copyright__ = "Copyright 2013 Colour Developers"
__license__ = "New BSD License - https://opensource.org/licenses/BSD-3-Clause"
__maintainer__ = "Colour Developers"
__email__ = "colour-developers@colour-science.org"
__status__ = "Production"

__all__ = [
    "TestOetf_BT601",
    "TestOetf_inverse_BT601",
]

class TestOetf_BT601(unittest.TestCase):
    """
    Define :func:`colour.models.rgb.transfer_functions.itur_bt_601.oetf_BT601`
    definition unit tests methods.
    """

    def test_oetf_BT601(self):
        """
        Test :func:`colour.models.rgb.transfer_functions.itur_bt_601.\
oetf_BT601` definition.
        """

        self.assertAlmostEqual(oetf_BT601(0.0), 0.0, places=7)

        self.assertAlmostEqual(oetf_BT601(0.015), 0.067500000000000, places=7)

        self.assertAlmostEqual(oetf_BT601(0.18), 0.409007728864150, places=7)

        self.assertAlmostEqual(oetf_BT601(1.0), 1.0, places=7)

    def test_n_dimensional_oetf_BT601(self):
        """
        Test :func:`colour.models.rgb.transfer_functions.itur_bt_601.\
oetf_BT601` definition n-dimensional arrays support.
        """

        L = 0.18
        E = oetf_BT601(L)

        L = np.tile(L, 6)
        E = np.tile(E, 6)
        np.testing.assert_almost_equal(oetf_BT601(L), E, decimal=7)

        L = np.reshape(L, (2, 3))
        E = np.reshape(E, (2, 3))
        np.testing.assert_almost_equal(oetf_BT601(L), E, decimal=7)

        L = np.reshape(L, (2, 3, 1))
        E = np.reshape(E, (2, 3, 1))
        np.testing.assert_almost_equal(oetf_BT601(L), E, decimal=7)

    def test_domain_range_scale_oetf_BT601(self):
        """
        Test :func:`colour.models.rgb.transfer_functions.itur_bt_601.\
oetf_BT601` definition domain and range scale support.
        """

        L = 0.18
        E = oetf_BT601(L)

        d_r = (("reference", 1), ("1", 1), ("100", 100))
        for scale, factor in d_r:
            with domain_range_scale(scale):
                np.testing.assert_almost_equal(
                    oetf_BT601(L * factor), E * factor, decimal=7
                )

    @ignore_numpy_errors
    def test_nan_oetf_BT601(self):
        """
        Test :func:`colour.models.rgb.transfer_functions.itur_bt_601.\
oetf_BT601` definition nan support.
        """

        oetf_BT601(np.array([-1.0, 0.0, 1.0, -np.inf, np.inf, np.nan]))

class TestOetf_inverse_BT601(unittest.TestCase):
    """
    Define :func:`colour.models.rgb.transfer_functions.itur_bt_601.\
oetf_inverse_BT601` definition unit tests methods.
    """

    def test_oetf_inverse_BT601(self):
        """
        Test :func:`colour.models.rgb.transfer_functions.itur_bt_601.\
oetf_inverse_BT601` definition.
        """

        self.assertAlmostEqual(oetf_inverse_BT601(0.0), 0.0, places=7)

        self.assertAlmostEqual(
            oetf_inverse_BT601(0.067500000000000), 0.015, places=7
        )

        self.assertAlmostEqual(
            oetf_inverse_BT601(0.409007728864150), 0.18, places=7
        )

        self.assertAlmostEqual(oetf_inverse_BT601(1.0), 1.0, places=7)

    def test_n_dimensional_oetf_inverse_BT601(self):
        """
        Test :func:`colour.models.rgb.transfer_functions.itur_bt_601.\
oetf_inverse_BT601` definition n-dimensional arrays support.
        """

        E = 0.409007728864150
        L = oetf_inverse_BT601(E)

        E = np.tile(E, 6)
        L = np.tile(L, 6)
        np.testing.assert_almost_equal(oetf_inverse_BT601(E), L, decimal=7)

        E = np.reshape(E, (2, 3))
        L = np.reshape(L, (2, 3))
        np.testing.assert_almost_equal(oetf_inverse_BT601(E), L, decimal=7)

        E = np.reshape(E, (2, 3, 1))
        L = np.reshape(L, (2, 3, 1))
        np.testing.assert_almost_equal(oetf_inverse_BT601(E), L, decimal=7)

    def test_domain_range_scale_oetf_inverse_BT601(self):
        """
        Test :func:`colour.models.rgb.transfer_functions.itur_bt_601.\
oetf_inverse_BT601` definition domain and range scale support.
        """

        E = 0.409007728864150
        L = oetf_inverse_BT601(E)

        d_r = (("reference", 1), ("1", 1), ("100", 100))
        for scale, factor in d_r:
            with domain_range_scale(scale):
                np.testing.assert_almost_equal(
                    oetf_inverse_BT601(E * factor), L * factor, decimal=7
                )

    @ignore_numpy_errors
    def test_nan_oetf_inverse_BT601(self):
        """
        Test :func:`colour.models.rgb.transfer_functions.itur_bt_601.\
oetf_inverse_BT601` definition nan support.
        """

        oetf_inverse_BT601(np.array([-1.0, 0.0, 1.0, -np.inf, np.inf, np.nan]))

if __name__ == "__main__":
    unittest.main()


# coding: utf-8

# ## 识别谷歌街景图片中的字母
#
# [street-view-getting-started-with-julia](https://www.kaggle.com/c/street-view-getting-started-with-julia) 让我们从谷歌街景的图片中鉴定字母，这个题目是让我们学习和使用Julia，Julia有python和R的易用性，有C语言的速度，无奈对Julia不是很熟悉，所以还是想用python来试试。

# In[1]:

import cv2
import numpy as np
import sys
import pandas as pd

# 我们希望所有的图片最后存储在一个numpy的矩阵当中，每一行为图片的像素值。为了得到统一的表达呢，我们将RGB三个通道的值做平均得到的灰度图像作为每个图片的表示:

# In[14]:

# typeData 为"train"或者"test"
# labelsInfo 包含每一个图片的ID
# 图片存储在trainResized和testResized文件夹内
def read_data(typeData, labelsInfo, imageSize):
    labelsIndex = labelsInfo["ID"]
    x = np.zeros((np.size(labelsIndex), imageSize))
    for idx, idImage in enumerate(labelsIndex):
        # 得到图片文件名并读取
        nameFile = typeData + "Resized/" + str(idImage) + ".Bmp"
        img = cv2.imread(nameFile)
        # 转化为灰度图
        temp = np.mean(img, 2)
        # 将图片转化为行向量
        x[idx, :] = np.reshape(temp, (1, imageSize))
    return x

# ### 预处理训练集和测试集

# In[15]:

imageSize = 400
trainlabels = pd.read_csv("trainLabels.csv")
testlabels = pd.read_csv("sampleSubmission.csv")
# 得到训练集的特征
xTrain = read_data('train', trainlabels, imageSize)
# 得到测试集的特征
xTest = read_data("test", testlabels, imageSize)

# #### 预览数据：

# In[19]:

print trainlabels.head(2)
print testlabels.head(2)

# In[20]:

yTrain = trainlabels["Class"]
yTrain = [ord(x) for x in yTrain]

# ## 模型训练
#
# ### 随机森林
#
# 使用随机森林进行训练，树的个数和深度需要多次调解寻求最佳值

# In[37]:

from sklearn.ensemble import RandomForestClassifier
get_ipython().magic(u'time rfc = RandomForestClassifier(n_estimators = 500, max_features = 50, max_depth=None)')
rfc.fit(xTrain, yTrain)

# #### 预测
# 将训练后的模型应用到测试集上，并保存结果：

# In[31]:

predTest = rfc.predict(xTest)
predResult = [chr(x) for x in predTest]
testlabels["Class"] = predResult
testlabels.to_csv("rf_500_50_result.csv",index = None)

# #### 结果
# 使用50颗树进行训练，提交kaggle之后准确率约为0.40
# 改用300颗树进行训练，提交kaggle之后准确率为0.46695
# 改用500颗树进行训练，深度为10，提价kaggle后准确率为0.40，估计出现了过拟合
# 改用500颗树进行训练，不设置深度，提价kaggle后准确率为0.47480

# ### 贝叶斯

# In[27]:

from sklearn.naive_bayes import GaussianNB as GNB
model_GNB = GNB()
model_GNB.fit(xTrain, yTrain)

predTest = model_GNB.predict(xTest)
predResult = [chr(x) for x in predTest]
testlabels["Class"] = predResult
testlabels.to_csv("gnb_result.csv",index = None)

# 贝叶斯的训练非常的快，把结果提交kaggle后，得到0.02389的准确率，明显低于随机森林

# ### GBDT

# In[36]:

from sklearn.ensemble import GradientBoostingClassifier
get_ipython().magic(u"time GBDT = GradientBoostingClassifier(loss='deviance', learning_rate=0.1, n_estimators=100, subsample=1.0,                         min_samples_split=2, min_samples_leaf=1, min_weight_fraction_leaf=0.0, max_depth=3, init=None,                         random_state=None, max_features=None, verbose=0, max_leaf_nodes=None, warm_start=False, presort='auto')")

get_ipython().magic(u'time GBDT.fit(xTrain, yTrain)')

get_ipython().magic(u'time predTest = GBDT.predict(xTest)')
predResult = [chr(x) for x in predTest]
testlabels["Class"] = predResult
testlabels.to_csv("gbdt_result.csv",index = None)

# 使用GBDT仅得到了0.31937的准确率，可能是我的默认参数没有调节好，关键是GBDT的训练时间太长，调试成本也比较高

# ### 神经网络

# In[40]:

import os
from skimage.io import imread
from lasagne import layers
from lasagne.nonlinearities import softmax
from nolearn.lasagne import NeuralNet, BatchIterator

# In[44]:

# Define functions
def read_datax(typeData, labelsInfo, imageSize, path):
    x = np.zeros((labelsInfo.shape[0], imageSize))

    for (index, idImage) in enumerate(labelsInfo['ID']):
        # use specially created 32 x 32 images
        nameFile = '{0}/{1}Resized32/{2}.Bmp'.format(path,
                    typeData, idImage)
        img = imread(nameFile, as_grey = True)

        x[index, :] = np.reshape(img, (1, imageSize))

    return x

def fit_model(reshaped_train_x, y, image_width,
                    image_height, reshaped_test_x):
    net = NeuralNet(
        layers = [
            ('input', layers.InputLayer),
            ('conv1', layers.Conv2DLayer),
            ('pool1', layers.MaxPool2DLayer),
            ('dropout1', layers.DropoutLayer),
            ('conv2', layers.Conv2DLayer),
            ('pool2', layers.MaxPool2DLayer),
            ('dropout2', layers.DropoutLayer),
            ('conv3', layers.Conv2DLayer),
            ('hidden4', layers.DenseLayer),
            ('output', layers.DenseLayer),
        ],
        input_shape = (None, 1, 32, 32),
        conv1_num_filters=32, conv1_filter_size=(5, 5),
        pool1_pool_size=(2, 2),
        dropout1_p=0.2,
        conv2_num_filters=64, conv2_filter_size=(5, 5),
        pool2_pool_size=(2, 2),
        dropout2_p=0.2,
        conv3_num_filters = 128, conv3_filter_size = (5, 5),
        hidden4_num_units=500,
        output_num_units = 62, output_nonlinearity = softmax,

        update_learning_rate = 0.01,
        update_momentum = 0.9,

        batch_iterator_train = BatchIterator(batch_size = 100),
        batch_iterator_test = BatchIterator(batch_size = 100),

        use_label_encoder = True,
        regression = False,
        max_epochs = 100,
        verbose = 1,
    )

    net.fit(reshaped_train_x, y)
    prediction = net.predict(reshaped_test_x)

    return prediction

# In[45]:

# 预处理数据，首先将图片保存为32*32的小图片
imageSize = 1024 # 32 x 32
image_width = image_height = int(imageSize ** 0.5)

labelsInfoTrain = pd.read_csv            ('trainLabels.csv'.format(path))
labelsInfoTest = pd.read_csv            ('sampleSubmission.csv'.format(path))

# Load dataset
nnxTrain = read_datax('train', labelsInfoTrain, imageSize, '.')
nnxTest = read_datax('test', labelsInfoTest, imageSize, '.')

nnyTrain = map(ord, labelsInfoTrain['Class'])
nnyTrain = np.array(yTrain)

# In[46]:

# 归一化数据
nnxTrain /= nnxTrain.std(axis = None)
nnxTrain -= nnxTrain.mean()

nnxTest /= nnxTest.std(axis = None)
nnxTest -= nnxTest.mean()

# In[47]:

# Reshape data
train_x_reshaped = nnxTrain.reshape(nnxTrain.shape[0], 1,
                  image_height, image_width).astype('float32')
test_x_reshaped = nnxTest.reshape(nnxTest.shape[0], 1,
                  image_height, image_width).astype('float32')

# In[54]:

# 进行训练和测试
predict = fit_model(train_x_reshaped, nnyTrain, image_width, image_height, test_x_reshaped)

# In[55]:

# 保存结果
yTest = map(chr, predict)
labelsInfoTest['Class'] = yTest
labelsInfoTest.to_csv('nnresult.csv'.format(path), index = False)

# 提交kaggle之后的准确率：0.64562