import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="1"


from pathlib import Path
from fastai.basic_train import load_learner
from fastai.text import load_data
from fastai.text.learner import language_model_learner
from fastai.text.models import AWD_LSTM, awd_lstm_lm_config

emb_sz=800
qrnn=False
bidir=False
n_layers=4
n_hid=2400


# https://app.wandb.ai/github/issues_lang_model/runs/22zkdqlr/overview
data_path = Path('/ds/lang_model')
model_path = data_path/'models_22zkdqlr'


def pass_through(x):
    return x


awd_lstm_lm_config.update(dict(emb_sz=emb_sz, qrnn=qrnn, bidir=bidir, n_layers=n_layers, n_hid=n_hid))

"""
Explanation: Location of Model Artifacts
Google Cloud Storage


model for inference (965 MB): https://storage.googleapis.com/issue_label_bot/model/lang_model/models_22zkdqlr/trained_model_22zkdqlr.hdf


encoder (for fine-tuning w/a classifier) (965 MB): 
https://storage.googleapis.com/issue_label_bot/model/lang_model/models_22zkdqlr/trained_model_encoder_22zkdqlr.pth


fastai.databunch (27.1 GB):
https://storage.googleapis.com/issue_label_bot/model/lang_model/data_save.hdf


checkpointed model (2.29 GB): 
https://storage.googleapis.com/issue_label_bot/model/lang_model/models_22zkdqlr/best_22zkdqlr.pth


Weights & Biases Run
https://app.wandb.ai/github/issues_lang_model/runs/22zkdqlr/overview
Part 1: Load Full Model + DataBunch In Order To Save Model For Inference
A fastai learner comes packaged with the training data and other data, however for inference we don't need this.  There is a way to export just the model weights without the data with learn.export or just the encoder base with learn.save_encoder.  Unfortunately, I forgot to do this during model training therefore we need to load the full checkpointed model and a databunch and save these artificacts for inference.
End of explanation
"""


data_lm = load_data(data_path, bs=96)

learn = language_model_learner(data=data_lm,
                               arch=AWD_LSTM,
                               model_dir=model_path,
                               pretrained=False)

"""
Explanation: Load learner object
Note: you don't have to do this over and over again, you just have to call learn.export() to save the learner after you have loaded everything.
End of explanation
"""


learn.load('best_22zkdqlr')

"""
Explanation: Load weights of trained model
End of explanation
"""


learn.export('trained_model_22zkdqlr.hdf')

learn.save_encoder('trained_model_encoder_22zkdqlr')

"""
Explanation: Export Minimal Model State For Inference
End of explanation
"""


del learn
del data_lm

"""
Explanation: The data is very large so if you are running this notebook best to release memory by deleting these objects and loading the more lightweight inference artifacts that we just saved.
End of explanation
"""


from inference import InferenceWrapper, pass_through

"""
Explanation: Part II:  Load Minimal Model For Inference
End of explanation
"""


wrapper = InferenceWrapper(model_path='/ds/lang_model/models_22zkdqlr/',
                           )

issue_string = '# hello abacadabra world \nA second line **something bold**.'

pooledfeat = wrapper.get_pooled_features(issue_string)
print(pooledfeat)
print(pooledfeat.shape)

rawfeat = wrapper.get_raw_features(issue_string)
print(rawfeat)
print(rawfeat.shape)

"""
Explanation: Create an InferenceWrapper object
End of explanation
"""


wrapper.learn.predict('I am having trouble opening a', 5)

"""
Explanation: Predict the next  5 words
We don't actually use this functionality, but it is interesting to see for those who are curious what the output of a langauge model is.  Recall that we are using the encoder of the language model to extract features from GitHub issues.
End of explanation
"""


