{"cells":[{"metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","jupyter":{"outputs_hidden":true},"trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python docker image: https://github.com/kaggle/docker-python\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport tensorflow as tf\nimport sys\nsys.path.extend([#'../input/tf2_0_baseline_w_bert/',#'../input/bert_modeling/',\n                 '../input/bert-joint-baseline/'])\nimport bert_utils\nimport modeling \n#import bert_optimization as optimization\nimport tokenization\nimport json\n\n#tf.compat.v1.disable_eager_execution()\n\n# Input data files are available in the \"../input/\" directory.\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n# In this case, we've got some extra BERT model files under `/kaggle/input/bertjointbaseline`\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# Any results you write to the current directory are saved as output.","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"on_kaggle_server = os.path.exists('/kaggle')\nnq_test_file = '../input/tensorflow2-question-answering/simplified-nq-test.jsonl' \npublic_dataset = os.path.getsize(nq_test_file)<20_000_000\nprivate_dataset = os.path.getsize(nq_test_file)>=20_000_000","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if True:\n    import importlib\n    importlib.reload(bert_utils)","execution_count":null,"outputs":[]},{"metadata":{"jupyter":{"outputs_hidden":true},"trusted":true},"cell_type":"code","source":"with open('../input/bert-joint-baseline/bert_config.json','r') as f:\n    config = json.load(f)\nprint(json.dumps(config,indent=4))\n","execution_count":null,"outputs":[]},{"metadata":{"jupyter":{"outputs_hidden":true},"trusted":true},"cell_type":"code","source":"class TDense(tf.keras.layers.Layer):\n    def __init__(self,\n                 output_size,\n                 kernel_initializer=None,\n                 bias_initializer=\"zeros\",\n                **kwargs):\n        super().__init__(**kwargs)\n        self.output_size = output_size\n        self.kernel_initializer = kernel_initializer\n        self.bias_initializer = bias_initializer\n    def build(self,input_shape):\n        dtype = tf.as_dtype(self.dtype or tf.keras.backend.floatx())\n        if not (dtype.is_floating or dtype.is_complex):\n          raise TypeError(\"Unable to build `TDense` layer with \"\n                          \"non-floating point (and non-complex) \"\n                          \"dtype %s\" % (dtype,))\n        input_shape = tf.TensorShape(input_shape)\n        if tf.compat.dimension_value(input_shape[-1]) is None:\n          raise ValueError(\"The last dimension of the inputs to \"\n                           \"`TDense` should be defined. \"\n                           \"Found `None`.\")\n        last_dim = tf.compat.dimension_value(input_shape[-1])\n        ### tf 2.1 rc min_ndim=3 -> min_ndim=2\n        self.input_spec = tf.keras.layers.InputSpec(min_ndim=2, axes={-1: last_dim})\n        self.kernel = self.add_weight(\n            \"kernel\",\n            shape=[self.output_size,last_dim],\n            initializer=self.kernel_initializer,\n            dtype=self.dtype,\n            trainable=True)\n        self.bias = self.add_weight(\n            \"bias\",\n            shape=[self.output_size],\n            initializer=self.bias_initializer,\n            dtype=self.dtype,\n            trainable=True)\n        super(TDense, self).build(input_shape)\n    def call(self,x):\n        return tf.matmul(x,self.kernel,transpose_b=True)+self.bias\n    \ndef mk_model(config):\n    seq_len = config['max_position_embeddings']\n    unique_id  = tf.keras.Input(shape=(1,),dtype=tf.int64,name='unique_id')\n    input_ids   = tf.keras.Input(shape=(seq_len,),dtype=tf.int32,name='input_ids')\n    input_mask  = tf.keras.Input(shape=(seq_len,),dtype=tf.int32,name='input_mask')\n    segment_ids = tf.keras.Input(shape=(seq_len,),dtype=tf.int32,name='segment_ids')\n    BERT = modeling.BertModel(config=config,name='bert')\n    pooled_output, sequence_output = BERT(input_word_ids=input_ids,\n                                          input_mask=input_mask,\n                                          input_type_ids=segment_ids)\n    \n    logits = TDense(2,name='logits')(sequence_output)\n    start_logits,end_logits = tf.split(logits,axis=-1,num_or_size_splits= 2,name='split')\n    start_logits = tf.squeeze(start_logits,axis=-1,name='start_squeeze')\n    end_logits   = tf.squeeze(end_logits,  axis=-1,name='end_squeeze')\n    \n    ans_type      = TDense(5,name='ans_type')(pooled_output)\n    return tf.keras.Model([input_ for input_ in [unique_id,input_ids,input_mask,segment_ids] \n                           if input_ is not None],\n                          [unique_id,start_logits,end_logits,ans_type],\n                          name='bert-baseline')    ","execution_count":null,"outputs":[]},{"metadata":{"jupyter":{"outputs_hidden":true},"trusted":true},"cell_type":"code","source":"small_config = config.copy()\nsmall_config['vocab_size']=16\nsmall_config['hidden_size']=64\nsmall_config['max_position_embeddings'] = 32\nsmall_config['num_hidden_layers'] = 4\nsmall_config['num_attention_heads'] = 4\nsmall_config['intermediate_size'] = 256\nsmall_config","execution_count":null,"outputs":[]},{"metadata":{"jupyter":{"outputs_hidden":true},"trusted":true},"cell_type":"code","source":"model= mk_model(config)","execution_count":null,"outputs":[]},{"metadata":{"jupyter":{"outputs_hidden":true},"trusted":true},"cell_type":"code","source":"model.summary()","execution_count":null,"outputs":[]},{"metadata":{"jupyter":{"outputs_hidden":true},"trusted":true},"cell_type":"code","source":"if False:\n    model_params = {v.name:v for v in model.trainable_variables}\n    model_roots = np.unique([v.name.split('/')[0] for v in model.trainable_variables])\n    print(model_roots)\n    saved_names = [k for k,v in tf.train.list_variables('../input/bertjointbaseline/bert_joint.ckpt')]\n    a_map = {v:v+':0' for v in saved_names}\n    model_roots = np.unique([v.name.split('/')[0] for v in model.trainable_variables])\n    def transform(x):\n        x = x.replace('attention/self','attention')\n        x = x.replace('attention','self_attention')\n        x = x.replace('attention/output','attention_output')  \n\n        x = x.replace('/dense','')\n        x = x.replace('/LayerNorm','_layer_norm')\n        x = x.replace('embeddings_layer_norm','embeddings/layer_norm')  \n\n        x = x.replace('attention_output_layer_norm','attention_layer_norm')  \n        x = x.replace('embeddings/word_embeddings','word_embeddings/embeddings')\n\n        x = x.replace('/embeddings/','/embedding_postprocessor/')  \n        x = x.replace('/token_type_embeddings','/type_embeddings')  \n        x = x.replace('/pooler/','/pooler_transform/')  \n        x = x.replace('answer_type_output_bias','ans_type/bias')  \n        x = x.replace('answer_type_output_','ans_type/')\n        x = x.replace('cls/nq/output_','logits/')\n        x = x.replace('/weights','/kernel')\n\n        return x\n    a_map = {k:model_params.get(transform(v),None) for k,v in a_map.items() if k!='global_step'}\n    tf.compat.v1.train.init_from_checkpoint(ckpt_dir_or_file='../input/bertjointbaseline/bert_joint.ckpt',\n                                            assignment_map=a_map)","execution_count":null,"outputs":[]},{"metadata":{"jupyter":{"outputs_hidden":true},"trusted":true},"cell_type":"code","source":"cpkt = tf.train.Checkpoint(model=model)\ncpkt.restore('../input/bert-joint-baseline/model_cpkt-1').assert_consumed()","execution_count":null,"outputs":[]},{"metadata":{"jupyter":{"outputs_hidden":true},"trusted":true},"cell_type":"code","source":"import tqdm\neval_records = \"../input/bert-joint-baseline/nq-test.tfrecords\"\n#nq_test_file = '../input/tensorflow2-question-answering/simplified-nq-test.jsonl'\nif on_kaggle_server and private_dataset:\n    eval_records='nq-test.tfrecords'\nif not os.path.exists(eval_records):\n    # tf2baseline.FLAGS.max_seq_length = 512\n    eval_writer = bert_utils.FeatureWriter(\n        filename=os.path.join(eval_records),\n        is_training=False)\n\n    tokenizer = tokenization.FullTokenizer(vocab_file='../input/bert-joint-baseline/vocab-nq.txt', \n                                           do_lower_case=True)\n\n    features = []\n    convert = bert_utils.ConvertExamples2Features(tokenizer=tokenizer,\n                                                   is_training=False,\n                                                   output_fn=eval_writer.process_feature,\n                                                   collect_stat=False)\n\n    n_examples = 0\n    tqdm_notebook= tqdm.tqdm_notebook if not on_kaggle_server else None\n    for examples in bert_utils.nq_examples_iter(input_file=nq_test_file, \n                                           is_training=False,\n                                           tqdm=tqdm_notebook):\n        for example in examples:\n            n_examples += convert(example)\n\n    eval_writer.close()\n    print('number of test examples: %d, written to file: %d' % (n_examples,eval_writer.num_features))","execution_count":null,"outputs":[]},{"metadata":{"jupyter":{"outputs_hidden":true},"trusted":true},"cell_type":"code","source":"seq_length = bert_utils.FLAGS.max_seq_length #config['max_position_embeddings']\nname_to_features = {\n      \"unique_id\": tf.io.FixedLenFeature([], tf.int64),\n      \"input_ids\": tf.io.FixedLenFeature([seq_length], tf.int64),\n      \"input_mask\": tf.io.FixedLenFeature([seq_length], tf.int64),\n      \"segment_ids\": tf.io.FixedLenFeature([seq_length], tf.int64),\n  }\n\ndef _decode_record(record, name_to_features=name_to_features):\n    \"\"\"Decodes a record to a TensorFlow example.\"\"\"\n    example = tf.io.parse_single_example(serialized=record, features=name_to_features)\n\n    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.\n    # So cast all int64 to int32.\n    for name in list(example.keys()):\n        t = example[name]\n        if name != 'unique_id': #t.dtype == tf.int64:\n            t = tf.cast(t, dtype=tf.int32)\n        example[name] = t\n\n    return example\n\ndef _decode_tokens(record):\n    return tf.io.parse_single_example(serialized=record, \n                                      features={\n                                          \"unique_id\": tf.io.FixedLenFeature([], tf.int64),\n                                          \"token_map\" :  tf.io.FixedLenFeature([seq_length], tf.int64)\n                                      })\n      \n","execution_count":null,"outputs":[]},{"metadata":{"jupyter":{"outputs_hidden":true},"trusted":true},"cell_type":"code","source":"raw_ds = tf.data.TFRecordDataset(eval_records)\ntoken_map_ds = raw_ds.map(_decode_tokens)\ndecoded_ds = raw_ds.map(_decode_record)\nds = decoded_ds.batch(batch_size=16,drop_remainder=False)","execution_count":null,"outputs":[]},{"metadata":{"jupyter":{"outputs_hidden":true},"trusted":false},"cell_type":"code","source":"# next(iter(decoded_ds))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"result=model.predict_generator(ds,verbose=1 if not on_kaggle_server else 0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"result[1]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"np.savez_compressed('bert-joint-baseline-output.npz',\n                    **dict(zip(['uniqe_id','start_logits','end_logits','answer_type_logits'],\n                               result)))","execution_count":null,"outputs":[]},{"metadata":{"jupyter":{"outputs_hidden":true},"trusted":true},"cell_type":"code","source":"def get_best_indexes(logits, n_best_size):\n  \"\"\"Get the n-best logits from a list.\"\"\"\n  index_and_score = sorted(\n      enumerate(logits[1:], 1), key=lambda x: x[1], reverse=True)\n  best_indexes = []\n  for i in range(len(index_and_score)):\n    if i >= n_best_size:\n      break\n    best_indexes.append(index_and_score[i][0])\n  return best_indexes\ndef top_k_indices(logits,n_best_size,token_map):\n    indices = np.argsort(logits[1:])+1\n    indices = indices[token_map[indices]!=-1]\n    return indices[-n_best_size:]\n    \n    \ndef compute_predictions(example):\n  \"\"\"Converts an example into an NQEval object for evaluation.\"\"\"\n  predictions = []\n  n_best_size = 10\n  max_answer_length = 30\n\n  for unique_id, result in example.results.items():\n    if unique_id not in example.features:\n      raise ValueError(\"No feature found with unique_id:\", unique_id)\n    token_map = np.array(example.features[unique_id][\"token_map\"]) #.int64_list.value\n    start_indexes = top_k_indices(result.start_logits,n_best_size,token_map)\n    if len(start_indexes)==0:\n        continue\n    end_indexes   = top_k_indices(result.end_logits,n_best_size,token_map)\n    if len(end_indexes)==0:\n        continue\n    indexes = np.array(list(np.broadcast(start_indexes[None],end_indexes[:,None])))  \n    indexes = indexes[(indexes[:,0]<indexes[:,1])*(indexes[:,1]-indexes[:,0]<max_answer_length)]\n    for i, (start_index,end_index) in enumerate(indexes):\n        summary = tf2baseline.ScoreSummary()\n        summary.short_span_score = (\n            result.start_logits[start_index] +\n            result.end_logits[end_index])\n        summary.cls_token_score = (\n            result.start_logits[0] + result.end_logits[0])\n        summary.answer_type_logits = result.answer_type_logits-result.answer_type_logits.mean()\n        start_span = token_map[start_index]\n        end_span = token_map[end_index] + 1\n\n        # Span logits minus the cls logits seems to be close to the best.\n        score = summary.short_span_score - summary.cls_token_score\n        predictions.append((score, i, summary, start_span, end_span))\n\n  # Default empty prediction.\n  score = -10000.0\n  short_span = tf2baseline.Span(-1, -1)\n  long_span = tf2baseline.Span(-1, -1)\n  summary = tf2baseline.ScoreSummary()\n\n  if predictions:\n    score, _, summary, start_span, end_span = sorted(predictions, reverse=True)[0]\n    short_span = tf2baseline.Span(start_span, end_span)\n    for c in example.candidates:\n      start = short_span.start_token_idx\n      end = short_span.end_token_idx\n      ## print(c['top_level'],c['start_token'],start,c['end_token'],end)\n      if c[\"top_level\"] and c[\"start_token\"] <= start and c[\"end_token\"] >= end:\n        long_span = tf2baseline.Span(c[\"start_token\"], c[\"end_token\"])\n        break\n\n  summary.predicted_label = {\n      \"example_id\": int(example.example_id),\n      \"long_answer\": {\n          \"start_token\": int(long_span.start_token_idx),\n          \"end_token\": int(long_span.end_token_idx),\n          \"start_byte\": -1,\n          \"end_byte\": -1\n      },\n      \"long_answer_score\": float(score),\n      \"short_answers\": [{\n          \"start_token\": int(short_span.start_token_idx),\n          \"end_token\": int(short_span.end_token_idx),\n          \"start_byte\": -1,\n          \"end_byte\": -1\n      }],\n      \"short_answer_score\": float(score),\n      \"yes_no_answer\": \"NONE\",\n      \"answer_type_logits\": summary.answer_type_logits.tolist(),\n      \"answer_type\": int(np.argmax(summary.answer_type_logits))\n  }\n\n  return summary\n\n\ndef compute_pred_dict(candidates_dict, dev_features, raw_results):\n    \"\"\"Computes official answer key from raw logits.\"\"\"\n    raw_results_by_id = [(int(res.unique_id),1, res) for res in raw_results]\n\n    examples_by_id = [(int(k),0,v) for k, v in candidates_dict.items()]\n  \n    features_by_id = [(int(d['unique_ids']),2,d) for d in dev_features] #list(zip(feature_ids, features))\n  \n    # Join examples with features and raw results.\n    examples = []\n    print('merging examples...')\n    merged = sorted(examples_by_id + raw_results_by_id + features_by_id)\n    print('done.')\n    for idx, type_, datum in merged:\n        if type_==0: #isinstance(datum, list):\n            examples.append(tf2baseline.EvalExample(idx, datum))\n        elif type_==2: #\"token_map\" in datum:\n            examples[-1].features[idx] = datum\n        else:\n            examples[-1].results[idx] = datum\n\n    # Construct prediction objects.\n    # tf.logging.info(\"Computing predictions...\")\n    print('Computing predictions...')\n    # summary_dict = {}\n    nq_pred_dict = {}\n    for e in tqdm.tqdm_notebook(examples):\n        summary = compute_predictions(e)\n        # summary_dict[e.example_id] = summary\n        nq_pred_dict[e.example_id] = summary.predicted_label\n\n    return nq_pred_dict","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nall_results = [bert_utils.RawResult(*x) for x in zip(*result)]\n    \nprint (\"Going to candidates file\")\n\ncandidates_dict = bert_utils.read_candidates('../input/tensorflow2-question-answering/simplified-nq-test.jsonl')\n\nprint (\"setting up eval features\")\n\neval_features = list(token_map_ds)\n\nprint (\"compute_pred_dict\")\n\ntqdm_notebook= tqdm.tqdm_notebook if not on_kaggle_server else None\nnq_pred_dict = bert_utils.compute_pred_dict(candidates_dict, \n                                       eval_features,\n                                       all_results,\n                                      tqdm=tqdm_notebook)\n\npredictions_json = {\"predictions\": list(nq_pred_dict.values())}\n\nprint (\"writing json\")\n\nwith tf.io.gfile.GFile('predictions.json', \"w\") as f:\n    json.dump(predictions_json, f, indent=4)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_short_answer(entry):\n    # if entry[\"short_answer_score\"] < 1.5:\n    #     return \"\"\n    \n    answer = []    \n    for short_answer in entry[\"short_answers\"]:\n        if short_answer[\"start_token\"] > -1:\n            answer.append(str(short_answer[\"start_token\"]) + \":\" + str(short_answer[\"end_token\"]))\n    if entry[\"yes_no_answer\"] != \"NONE\":\n        answer.append(entry[\"yes_no_answer\"])\n    return \" \".join(answer)\n\ndef create_long_answer(entry):\n   # if entry[\"long_answer_score\"] < 1.5:\n   # return \"\"\n\n    answer = []\n    if entry[\"long_answer\"][\"start_token\"] > -1:\n        answer.append(str(entry[\"long_answer\"][\"start_token\"]) + \":\" + str(entry[\"long_answer\"][\"end_token\"]))\n    return \" \".join(answer)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_answers_df = pd.read_json(\"../working/predictions.json\")\nfor var_name in ['long_answer_score','short_answer_score','answer_type']:\n    test_answers_df[var_name] = test_answers_df['predictions'].apply(lambda q: q[var_name])\ntest_answers_df[\"long_answer\"] = test_answers_df[\"predictions\"].apply(create_long_answer)\ntest_answers_df[\"short_answer\"] = test_answers_df[\"predictions\"].apply(create_short_answer)\ntest_answers_df[\"example_id\"] = test_answers_df[\"predictions\"].apply(lambda q: str(q[\"example_id\"]))\n\nlong_answers = dict(zip(test_answers_df[\"example_id\"], test_answers_df[\"long_answer\"]))\nshort_answers = dict(zip(test_answers_df[\"example_id\"], test_answers_df[\"short_answer\"]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_submission = pd.read_csv(\"../input/tensorflow2-question-answering/sample_submission.csv\")\n\nlong_prediction_strings = sample_submission[sample_submission[\"example_id\"].str.contains(\"_long\")].apply(lambda q: long_answers[q[\"example_id\"].replace(\"_long\", \"\")], axis=1)\nshort_prediction_strings = sample_submission[sample_submission[\"example_id\"].str.contains(\"_short\")].apply(lambda q: short_answers[q[\"example_id\"].replace(\"_short\", \"\")], axis=1)\n\nsample_submission.loc[sample_submission[\"example_id\"].str.contains(\"_long\"), \"PredictionString\"] = long_prediction_strings\nsample_submission.loc[sample_submission[\"example_id\"].str.contains(\"_short\"), \"PredictionString\"] = short_prediction_strings\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sample_submission.to_csv(\"submission.csv\", index=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"if public_dataset:\n    print(test_answers_df[\"long_answer_score\"].describe())","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"if public_dataset:\n    print(np.bincount(test_answers_df['answer_type'].values))","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"if public_dataset:\n    print(test_answers_df[test_answers_df['answer_type']==0])","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"if public_dataset:\n    print(test_answers_df.predictions.values[-4])","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"if public_dataset:\n    print(sample_submission.head())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class ShowPrediction:\n    def __init__(self,jsonl_file):\n        self._data = {}\n        with open(jsonl_file,'r') as f:\n            for line in f.readlines():\n                d = json.loads(line)\n                #print(d.keys())\n                self._data[int(d['example_id'])]={\n                    'text': d['document_text'],\n                    'question': d['question_text']\n                }\n    def __call__(self,prediction,include_full_text=True):\n        data = self._data[prediction['example_id']]\n        res = {'question': data['question']}\n        if include_full_text:\n            res['text'] = data['text']\n        for type_ in ['long_answer','short_answers']:\n            ans = prediction[type_]\n            if isinstance(ans,list):\n                ans = ans[0]\n            start,end = ans['start_token'],ans['end_token']\n            res[type_] = ' '.join(data['text'].split()[start:end])\n        return res\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"if public_dataset:\n    show_pred = ShowPrediction('../input/tensorflow2-question-answering/simplified-nq-test.jsonl')","execution_count":null,"outputs":[]},{"metadata":{"jupyter":{"outputs_hidden":true},"trusted":false},"cell_type":"code","source":"if public_dataset:\n    for pred in test_answers_df.predictions[test_answers_df.answer_type==0]:\n        print(json.dumps(show_pred(pred,include_full_text=True),indent=4))","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"if public_dataset:\n    for pred in np.random.choice(predictions_json['predictions'],10):\n        print(json.dumps(show_pred(pred,include_full_text=False),indent=4))","execution_count":null,"outputs":[]},{"metadata":{"trusted":false},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.6.5"}},"nbformat":4,"nbformat_minor":1}