syntree.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nltk\n",
    "from nltk.corpus import wordnet as wn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import spacy\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "def make_to_path(path):\n",
    "    return Path(path) if(type(path) == str) else path\n",
    "def remove_if_exists(path):\n",
    "    path = make_to_path(path)\n",
    "    if(path.exists() and path.is_dir()):\n",
    "        shutil.rmtree(path)\n",
    "    elif(path.exists()):\n",
    "        path.unlink()\n",
    "def create_dir(path):\n",
    "    path = make_to_path(path)\n",
    "    if(not path.exists()):  path.mkdir(parents=True, exist_ok=True)\n",
    "    return path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "remove_if_exists('./data')\n",
    "create_dir('./data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create bs batches, tot_len//bs seqlen <- in each file\n",
    "class TextGenerator():\n",
    "    def __init__(self, path_list=None, text_list=None, dest_path=Path('./data/text_gen_data'), nlp=None, encoding=\"utf8\"):\n",
    "        assert type(path_list) == list or type(text_list) == list\n",
    "        dest_path = make_to_path(dest_path)\n",
    "        dest_path = create_dir(dest_path)\n",
    "        if(path_list is None):\n",
    "            self.path_list = self.create_data_paths(dest_path, text_list)\n",
    "        else:\n",
    "            self.path_list = path_list\n",
    "        self.nlp = spacy.load(\"en_core_web_sm\") if(nlp is None) else nlp\n",
    "        self.encoding = encoding\n",
    "    \n",
    "    def create_data_paths(self, base_path, text_list):\n",
    "        i = 0\n",
    "        paths = []\n",
    "        for text in text_list:\n",
    "            dest_path = base_path / f\"{i}.txt\"\n",
    "            i += 1\n",
    "            paths.append(dest_path)\n",
    "            with open(dest_path, \"w\") as f:\n",
    "                f.write(text)\n",
    "        return paths\n",
    "    \n",
    "    def get_next_text(self):\n",
    "        for p in self.path_list:\n",
    "            with open(p, 'r',encoding=self.encoding) as f:\n",
    "                s = f.read()\n",
    "                yield p, s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sentences = [\n",
    "    'Hariom is playing',\n",
    "    'I got milk from the store',\n",
    "    'from the store, I got milk',\n",
    "    'hey hey hey',\n",
    "    'this sentence can be structured to a tree in a systematic manner',\n",
    "    'what is he doing',\n",
    "    'sentence illegal is',\n",
    "    'sentence is legal',\n",
    "    'This sentence is going to be long because we need to check what is going on',\n",
    "    'The poop is kept on the beautiful table',\n",
    "    'bread and butter is good for health',\n",
    "    'wine is great with bread and butter is good for health'\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlp = spacy.load(\"en_core_web_sm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Vocab():\n",
    "    def __init__(self, tokenizer, max_len=20000, nlp=None):\n",
    "        self.tokenizer, self.max_len = tokenizer, max_len\n",
    "        self.nlp = spacy.load(\"en_core_web_sm\") if(nlp is None) else nlp\n",
    "        self.itos, self.stoi, self.stoc = [], {}, {}\n",
    "        self.special_tokens = [\"xxbos\", \"xxeos\", \"xxmaj\", \"xxcap\", \"xxunk\"]\n",
    "        \n",
    "    def trim_word_index_mappings(self):\n",
    "        sorted_stoc = {k: v for k, v in reversed(sorted(self.stoc.items(), key=lambda item: item[1]))}\n",
    "        for k in self.special_tokens:\n",
    "            if k in sorted_stoc: del sorted_stoc[k]\n",
    "        sorted_list = list(sorted_stoc.items())\n",
    "        \n",
    "        new_stoi, new_itos, new_stoc = {}, [], {}\n",
    "        for spcl_tok in self.special_tokens:\n",
    "            cur_ind = len(new_itos)\n",
    "            new_stoi[spcl_tok] = cur_ind\n",
    "            new_itos.append(spcl_tok)\n",
    "            \n",
    "        for i in range(min(self.max_len, len(sorted_list))):\n",
    "            k, v = sorted_list[i]\n",
    "            cur_ind = len(new_itos)\n",
    "            new_stoc[k], new_stoi[k] = v, cur_ind\n",
    "            new_itos.append(k)\n",
    "        self.itos, self.stoi, self.stoc = new_itos, new_stoi, new_stoc\n",
    "    \n",
    "    def update_token_list(self, token_list, word_to_leaf_type):\n",
    "        new_list = []\n",
    "        for token in token_list:\n",
    "            tok, ltype = token\n",
    "            if tok not in self.stoi: new_list.append((\"xxunk\", \"SPCL_APPEND\"))\n",
    "            else: new_list.append((tok, ltype))\n",
    "        return new_list\n",
    "    \n",
    "    def push_tok_to_word_index_mappings(self, tok):\n",
    "        if tok not in self.stoi:\n",
    "            ind = len(self.itos)\n",
    "            self.stoi[tok], self.stoc[tok] = ind, 1\n",
    "            self.itos.append(tok)\n",
    "        else:\n",
    "            self.stoc[tok] += 1\n",
    "        \n",
    "    def get_trimmed_tokens(self, token_list, word_to_leaf_type):\n",
    "        return self.update_token_list(token_list, word_to_leaf_type)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Tokenizer():\n",
    "    def __init__(self, path_list=None, text_list=None, dest_path=Path('./data/tokenizer_data'), \n",
    "                 apply_defaults=True, max_len=2000, callbacks=[], \n",
    "                 nlp=None):\n",
    "        dest_path = make_to_path(dest_path)\n",
    "        dest_path = create_dir(dest_path)\n",
    "        self.leaf_types = {\n",
    "            \"ROOT_LEAF\": 0,\n",
    "            \"ONE_CHILD_LEAF\": 1,\n",
    "            \"TWO_CHILD_LEAF\": 2,\n",
    "            \"BRIDGE_LEAF\": 3,\n",
    "            \"INVARIANT\": 4,\n",
    "        }       \n",
    "        self.word_to_leaf_type = {\n",
    "            \"PUNCT\": self.leaf_types[\"BRIDGE_LEAF\"],\n",
    "            \"SYM\": self.leaf_types[\"ONE_CHILD_LEAF\"],\n",
    "            \"ADJ\": self.leaf_types[\"ONE_CHILD_LEAF\"],\n",
    "            \"ADP\": self.leaf_types[\"ONE_CHILD_LEAF\"],\n",
    "            \"CCONJ\": self.leaf_types[\"BRIDGE_LEAF\"],\n",
    "            \"NUM\": self.leaf_types[\"ONE_CHILD_LEAF\"],\n",
    "            \"DET\": self.leaf_types[\"ONE_CHILD_LEAF\"],\n",
    "            \"ADV\": self.leaf_types[\"TWO_CHILD_LEAF\"],\n",
    "            \n",
    "            \"X\": self.leaf_types[\"ONE_CHILD_LEAF\"],\n",
    "            \"VERB\": self.leaf_types[\"TWO_CHILD_LEAF\"],\n",
    "            \"NOUN\": self.leaf_types[\"ONE_CHILD_LEAF\"],\n",
    "            \"PROPN\": self.leaf_types[\"ONE_CHILD_LEAF\"],\n",
    "            \"PART\": self.leaf_types[\"ONE_CHILD_LEAF\"],\n",
    "            \"INTJ\": self.leaf_types[\"BRIDGE_LEAF\"],\n",
    "            \"PRON\": self.leaf_types[\"ONE_CHILD_LEAF\"],\n",
    "            \"SPCL_APPEND\": self.leaf_types[\"INVARIANT\"],\n",
    "            \"SPCL_BRIDGE\": self.leaf_types[\"BRIDGE_LEAF\"],\n",
    "            \"ROOT\": self.leaf_types[\"ROOT_LEAF\"]\n",
    "        }\n",
    "        \n",
    "        self.nlp = spacy.load(\"en_core_web_sm\") if(nlp is None) else nlp\n",
    "        self.text_generator = TextGenerator(path_list, text_list, nlp=self.nlp)\n",
    "        self.tokens_paths = []\n",
    "        self.callbacks = [self.put_bos_eos, self.deal_with_caps] if(apply_defaults) else []\n",
    "        self.callbacks.extend(callbacks)\n",
    "        self.tokens_len = self.prepare_initial_tokens(path_list, dest_path)\n",
    "        self.max_len = max_len\n",
    "        self.vocab = Vocab(self, max_len, nlp=self.nlp)\n",
    "        self.prepare_vocab()\n",
    "        self.update_tokens_lists()\n",
    "        \n",
    "    def tokenify(self, vector):\n",
    "        lv = [i.item() for i in vector]\n",
    "        ret = []\n",
    "        for item in lv:\n",
    "            ret.append(tk.vocab.itos[item])\n",
    "        return ret\n",
    "    \n",
    "    def put_bos_eos(self, tokens_list):\n",
    "        return [(\"xxbos\", \"SPCL_BRIDGE\"), *tokens_list, (\"xxeos\", \"SPCL_BRIDGE\")]\n",
    "\n",
    "    def deal_with_caps(self, token_list):\n",
    "        nlist = []\n",
    "        for token in token_list:\n",
    "            nlist.append((token[0].lower(), token[1]))\n",
    "            if(token[0].isupper()):\n",
    "                nlist.append((\"xxcap\", \"SPCL_APPEND\"))\n",
    "            elif(token[0][0].isupper()):\n",
    "                nlist.append((\"xxmaj\", \"SPCL_APPEND\"))\n",
    "        return nlist\n",
    "\n",
    "    def apply_callbacks(self, token_list):\n",
    "        for callback in self.callbacks:\n",
    "            token_list = callback(token_list)\n",
    "        return token_list\n",
    "    \n",
    "    def reset_vocab(self):\n",
    "        self.vocab = Vocab(self, self.max_len)\n",
    "\n",
    "    def prepare_vocab(self):\n",
    "        for _,token_list in self.get_next_token_list():\n",
    "            for tok,_ in token_list:\n",
    "                self.vocab.push_tok_to_word_index_mappings(tok)\n",
    "        self.vocab.trim_word_index_mappings()\n",
    "        \n",
    "    def prepare_initial_tokens(self, src_paths, dest_root):\n",
    "        tot_len = 0\n",
    "        for path, content in self.text_generator.get_next_text():\n",
    "            fname = path.stem\n",
    "            dest_path = dest_root / f\"{fname}.pkl\"\n",
    "            self.tokens_paths.append(dest_path)\n",
    "            token_list = self.simplified_tagged_list(content)\n",
    "            token_list = self.apply_callbacks(token_list)\n",
    "            tot_len += len(token_list)\n",
    "            with open(dest_path, 'wb') as f:\n",
    "                pickle.dump(token_list, f)\n",
    "        return tot_len\n",
    "    \n",
    "    def get_next_token_list(self):\n",
    "        for p in self.tokens_paths:\n",
    "            with open(p, \"rb\") as f:\n",
    "                tk_list = pickle.load(f)\n",
    "                for i in range(0, len(tk_list), self.max_len):\n",
    "                    yield p, tk_list[i : min(i+self.max_len, len(tk_list))]\n",
    "\n",
    "    def update_tokens_lists(self):\n",
    "        for p in self.tokens_paths:\n",
    "            with open(p, \"rb\") as f:\n",
    "                tk_list = pickle.load(f)\n",
    "                new_tk_list = self.vocab.get_trimmed_tokens(tk_list, self.word_to_leaf_type)\n",
    "            with open(p, \"wb\") as f:\n",
    "                pickle.dump(new_tk_list, f)\n",
    "    \n",
    "    def simplified_tagged_list(self, sentence):\n",
    "        doc = self.nlp(sentence)\n",
    "        tags = [(token.text, token.pos_) for token in doc]\n",
    "        return tags\n",
    "\n",
    "    def get_indices_from_tokens(self, tokens):\n",
    "        return [self.vocab.stoi[token[0]] for token in tokens]\n",
    "    \n",
    "    def get_word_tensors(self, tokens):\n",
    "        indices = self.get_indices_from_tokens(tokens)\n",
    "        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "        return torch.tensor(indices, dtype=torch.long, device=device).view(-1, 1)\n",
    "    def get_ltype(self, token):\n",
    "        return self.word_to_leaf_type[token]\n",
    "    def get_ltypes_dict(self):\n",
    "        return self.leaf_types\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_brown_text():\n",
    "    from nltk.corpus import brown\n",
    "    return brown.words(categories='news')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "btext_list = get_brown_text()\n",
    "btext = \"\"\n",
    "for word in btext_list:\n",
    "    btext += f\" {word}\"\n",
    "btext_list = [btext]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nltk.tag import pos_tag, map_tag\n",
    "from nltk.tokenize import word_tokenize\n",
    "\n",
    "def unsimplified_tagged_list(sentence):\n",
    "    tokens = word_tokenize(sentence)\n",
    "    return pos_tag(tokens)\n",
    "    \n",
    "def tagged_list(sentence):\n",
    "    tokens = word_tokenize(sentence)\n",
    "    posTagged = pos_tag(tokens)\n",
    "    tags = [(word, map_tag('en-ptb', 'universal', tag)) for word, tag in posTagged]\n",
    "    return tags\n",
    "\n",
    "def download_tag_packs():\n",
    "    import nltk\n",
    "    nltk.download('punkt')\n",
    "    nltk.download([\"tagsets\", \"universal_tagset\"])\n",
    "    nltk.download(\"maxent_treebank_pos_tagger\")\n",
    "    nltk.download(\"maxent_ne_chunker\")\n",
    "    nltk.download(\"punkt\")\n",
    "    nltk.download('wordnet')\n",
    "    nltk.download('averaged_perceptron_tagger')\n",
    "    nltk.download('brown')\n",
    "\n",
    "def get_brown_text():\n",
    "    from nltk.corpus import brown\n",
    "    return brown.words(categories='news')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class InvalidDirectionError(Exception):\n",
    "    def __init__(self, *args):\n",
    "        if args:\n",
    "            self.message = args[0]\n",
    "        else:\n",
    "            self.message = None\n",
    "\n",
    "    def __str__(self):\n",
    "        if self.message:\n",
    "            return f'InvalidDirectionError: {self.message}'\n",
    "        else:\n",
    "            return 'InvalidDirectionError: '"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BaseLeaf():\n",
    "    def __init__(self, token, wtype, tokenizer, embedding_size=50):\n",
    "        self.tokenizer = tokenizer\n",
    "        self.tokens, self.leaf_type = [(token, wtype)], tokenizer.get_ltype(wtype)\n",
    "        self.leaf_types_dict = self.tokenizer.get_ltypes_dict()\n",
    "        self.example_tokens, self.target_tokens, self.example_ltypes, self.target_ltypes = None, None, None, None\n",
    "\n",
    "    def label_tokens(self, example_tokens, target_tokens):\n",
    "        self.example_tokens, self.target_tokens = ((self.tokenizer.get_word_tensors(example_tokens), self.tokenizer.get_word_tensors(target_tokens))\n",
    "                                                  if self.get_self_leaf_type() is not self.leaf_types_dict[\"ROOT_LEAF\"]\n",
    "                                                  else (None, None))\n",
    "    \n",
    "    def label_ltypes(self, example_ltypes, target_ltypes):\n",
    "        self.example_ltypes, self.target_ltypes = ((torch.tensor(example_ltypes).view(-1,1), torch.tensor(target_ltypes).view(-1,1))\n",
    "                                                  if self.get_self_leaf_type() is not self.leaf_types_dict[\"ROOT_LEAF\"]\n",
    "                                                  else (None, None))\n",
    "    \n",
    "    def get_labelled_set(self):\n",
    "        return self.example_tokens, self.target_tokens, self.example_ltypes, self.target_ltypes\n",
    "\n",
    "    def append_token(self, token, wtype):\n",
    "        self.tokens.append((token, wtype))\n",
    "        \n",
    "    def get_self_leaf_type(self):\n",
    "        return self.leaf_type\n",
    "    def generate_text(self):\n",
    "        before, after = self.get_ordered_children()\n",
    "        s = \"\"\n",
    "        for c in before:\n",
    "            s += c.generate_text() if(c is not None) else \"\"\n",
    "        for token, wtype in self.tokens:\n",
    "            s += f\" {token}\"\n",
    "        for c in after:\n",
    "            s += c.generate_text() if(c is not None) else \"\"\n",
    "        return s\n",
    "    \n",
    "    def get_next_leaf(self):\n",
    "        before, after = self.get_ordered_children()\n",
    "        for child in before:\n",
    "            yield from child.get_next_leaf()\n",
    "        yield self\n",
    "        for child in after:\n",
    "            yield from child.get_next_leaf()\n",
    "            \n",
    "    def get_leaf_type_of(self, wtype):\n",
    "        return self.tokenizer.get_ltype(wtype)\n",
    "    def __str__(self, spacing):\n",
    "        raise NotImplementedError()\n",
    "    def push_token(self, token, wtype):\n",
    "        raise NotImplementedError()\n",
    "    def set_parent(self, leaf):\n",
    "        raise NotImplementedError()\n",
    "    def replace_child(self, old, new):\n",
    "        raise NotImplementedError()\n",
    "    def get_ordered_children(self):\n",
    "        raise NotImplementedError()\n",
    "    def get_ordered_children_for_model(self):\n",
    "        raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TwoChildLeaf(BaseLeaf):\n",
    "    def __init__(self, token, wtype, tokenizer, creator=None):\n",
    "        super().__init__(token, wtype, tokenizer)\n",
    "        assert (self.get_self_leaf_type() == self.leaf_types_dict[\"TWO_CHILD_LEAF\"] or\n",
    "               self.get_self_leaf_type() == self.leaf_types_dict[\"INVARIANT\"])\n",
    "        self.l_child, self.r_child = None, None\n",
    "        self.parent = creator\n",
    "        self.bridge = None\n",
    "    \n",
    "    def get_ordered_children(self):\n",
    "        before, after = [], []\n",
    "        if(self.l_child is not None): before.append(self.l_child)\n",
    "        if(self.r_child is not None): after.append(self.r_child)\n",
    "        if(self.bridge is not None): after.append(self.bridge)\n",
    "        return before, after\n",
    "    \n",
    "    def get_ordered_children_for_model(self):\n",
    "        before, after = [], []\n",
    "        if(self.l_child is not None): before.append(self.l_child)\n",
    "        if(self.r_child is not None): before.append(self.r_child)\n",
    "        if(self.bridge is not None): after.append(self.bridge)\n",
    "        return before, after\n",
    "        \n",
    "    def __str__(self, spacing):\n",
    "        new_spacing = f\"{spacing}\\t\"\n",
    "        base_str = f\"{spacing}type:{self.get_self_leaf_type()} tokens:{self.tokens}\\n\"\n",
    "        if(self.l_child is not None):\n",
    "            base_str += f\"{self.l_child.__str__(new_spacing)}\\n\"\n",
    "        else:\n",
    "            base_str += f\"{new_spacing}None\\n\"\n",
    "        if(self.r_child is not None):\n",
    "            base_str += f\"{self.r_child.__str__(new_spacing)}\\n\"\n",
    "        else:\n",
    "            base_str += f\"{new_spacing}None\\n\"\n",
    "        if(self.bridge is not None):\n",
    "            base_str += f\"{self.bridge.__str__(new_spacing)}\\n\"\n",
    "        else:\n",
    "            base_str += f\"{new_spacing}None\\n\"\n",
    "        return base_str\n",
    "\n",
    "    def set_parent(self, parent):\n",
    "        self.parent = parent\n",
    "    \n",
    "    def replace_child(self, old, new):\n",
    "        if(self.l_child == old):\n",
    "            self.l_child = new\n",
    "        elif(self.r_child == old):\n",
    "            self.r_child = new\n",
    "        else:\n",
    "            raise Exception('does not exist as child')\n",
    "            \n",
    "    def push_token(self, token, wtype):\n",
    "        ltype = self.get_leaf_type_of(wtype)\n",
    "        if(ltype == self.leaf_types_dict[\"BRIDGE_LEAF\"]):\n",
    "            self.bridge = BridgeLeaf(token, wtype, self.tokenizer, self)\n",
    "            return self.bridge\n",
    "        elif(ltype == self.leaf_types_dict[\"TWO_CHILD_LEAF\"] or ltype == self.leaf_types_dict[\"INVARIANT\"]):\n",
    "            self.append_token(token, wtype)\n",
    "            return self\n",
    "        elif(ltype == self.leaf_types_dict[\"ONE_CHILD_LEAF\"]):\n",
    "            self.r_child = OneChildLeaf(token, wtype, self.tokenizer, self)\n",
    "            return self.r_child\n",
    "        else:\n",
    "            raise Exception('wrong leaf type for child')\n",
    "        \n",
    "class OneChildLeaf(BaseLeaf):\n",
    "    def __init__(self, token, wtype, tokenizer, creator=None):\n",
    "        super().__init__(token, wtype, tokenizer)\n",
    "        assert (self.get_self_leaf_type() == self.leaf_types_dict[\"ONE_CHILD_LEAF\"] or\n",
    "               self.get_self_leaf_type() == self.leaf_types_dict[\"INVARIANT\"])\n",
    "        self.bridge, self.child = None, None\n",
    "        self.parent = creator\n",
    "    \n",
    "    def get_ordered_children(self):\n",
    "        before = []\n",
    "        after = []\n",
    "        if(self.child is not None): after.append(self.child)\n",
    "        if(self.bridge is not None): after.append(self.bridge)\n",
    "        return before, after\n",
    "        \n",
    "    def get_ordered_children_for_model(self):\n",
    "        before, after = [], []\n",
    "        if(self.child is not None): before.append(self.child)\n",
    "        if(self.bridge is not None): after.append(self.bridge)\n",
    "        return before, after    \n",
    "            \n",
    "    def __str__(self, spacing):\n",
    "        new_spacing = f\"{spacing}\\t\"\n",
    "        base_str = f\"{spacing}type:{self.get_self_leaf_type()} words:{self.tokens}\\n\"\n",
    "        if(self.child is not None):\n",
    "            base_str += f\"{self.child.__str__(new_spacing)}\\n\"\n",
    "        else:\n",
    "            base_str += f\"{new_spacing}None\\n\"\n",
    "        \n",
    "        if(self.bridge is not None):\n",
    "            base_str += f\"{self.bridge.__str__(new_spacing)}\\n\"\n",
    "        else:\n",
    "            base_str += f\"{new_spacing}None\\n\"\n",
    "        return base_str\n",
    "\n",
    "    def set_parent(self, parent):\n",
    "        self.parent = parent\n",
    "    \n",
    "    def replace_child(self, old, new):\n",
    "        if(self.child == old):\n",
    "            self.child = new\n",
    "        else:\n",
    "            raise Exception('does not exist as child')\n",
    "    \n",
    "    def push_token(self, token, wtype):\n",
    "        ltype = self.get_leaf_type_of(wtype)\n",
    "        if(ltype == self.leaf_types_dict[\"BRIDGE_LEAF\"]):\n",
    "            self.bridge = BridgeLeaf(token, wtype, self.tokenizer, self)\n",
    "            return self.bridge\n",
    "        elif(ltype == self.leaf_types_dict[\"ONE_CHILD_LEAF\"] or ltype == self.leaf_types_dict[\"INVARIANT\"]):\n",
    "            self.append_token(token, wtype)\n",
    "            return self\n",
    "        elif(ltype == self.leaf_types_dict[\"TWO_CHILD_LEAF\"]):\n",
    "            if(self.parent.get_self_leaf_type() == self.leaf_types_dict[\"ONE_CHILD_LEAF\"]):\n",
    "                self.parent.push_token(token, wtype)\n",
    "            else:\n",
    "                # replace new_leaf as this node\n",
    "                new_leaf = TwoChildLeaf(token, wtype, self.tokenizer, None)\n",
    "                new_leaf.l_child = self\n",
    "                if(self.parent is not None):\n",
    "                    self.parent.replace_child(self, new_leaf)   # now new_leaf is the child of the parent\n",
    "                    new_leaf.set_parent(self.parent)\n",
    "                    self.set_parent(new_leaf)\n",
    "                return new_leaf\n",
    "        else:\n",
    "            raise Exception('wrong leaf type for child')\n",
    "            \n",
    "class BridgeLeaf(BaseLeaf):\n",
    "    def __init__(self, token, wtype, tokenizer, creator=None):\n",
    "        super().__init__(token, wtype, tokenizer)\n",
    "        assert (self.get_self_leaf_type() == self.leaf_types_dict[\"BRIDGE_LEAF\"] or\n",
    "               self.get_self_leaf_type() == self.leaf_types_dict[\"INVARIANT\"])\n",
    "        self.parent = creator\n",
    "        self.child = None\n",
    "    \n",
    "    def get_ordered_children(self):\n",
    "        before, after = [], []\n",
    "        if(self.child is not None): after.append(self.child)\n",
    "        return before, after\n",
    "    \n",
    "    def get_ordered_children_for_model(self):\n",
    "        before, after = [], []\n",
    "        if(self.child is not None): after.append(self.child)\n",
    "        return before, after\n",
    "    \n",
    "    def __str__(self, spacing):\n",
    "        new_spacing = f\"{spacing}\\t\"\n",
    "        base_str = f\"{spacing}type:{self.get_self_leaf_type()} words:{self.tokens}\\n\"\n",
    "        if(self.child is not None):\n",
    "            base_str += f\"{self.child.__str__(new_spacing)}\\n\"\n",
    "        else:\n",
    "            base_str += f\"{new_spacing}None\\n\"\n",
    "        \n",
    "        return base_str\n",
    "    \n",
    "    def set_parent(self, parent):\n",
    "        self.parent = parent\n",
    "        \n",
    "    def replace_child(self, old, new):\n",
    "        if(self.child == old):\n",
    "            self.child = new\n",
    "        else:\n",
    "            raise Exception('does not exist as child')\n",
    "    \n",
    "    def push_token(self, token, wtype):\n",
    "        ltype = self.get_leaf_type_of(wtype)\n",
    "        if(ltype == self.leaf_types_dict[\"BRIDGE_LEAF\"] or ltype == self.leaf_types_dict[\"INVARIANT\"]):\n",
    "            self.append_token(token, wtype)\n",
    "            return self\n",
    "        else:\n",
    "            if(ltype == self.leaf_types_dict[\"TWO_CHILD_LEAF\"]):\n",
    "                new_leaf = TwoChildLeaf(token, wtype, self.tokenizer, self)\n",
    "            elif(ltype == self.leaf_types_dict[\"ONE_CHILD_LEAF\"]):\n",
    "                new_leaf = OneChildLeaf(token, wtype, self.tokenizer, self)\n",
    "            else:\n",
    "                raise Exception('wrong leaf type for child')\n",
    "            self.child = new_leaf\n",
    "            return self.child\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RootLeaf(BaseLeaf):\n",
    "    def __init__(self, tokenizer):\n",
    "        super().__init__(None, \"ROOT\", tokenizer)\n",
    "        assert self.get_self_leaf_type() == self.leaf_types_dict[\"ROOT_LEAF\"]\n",
    "        self.child = None\n",
    "        self.example_tokens, self.target_tokens, self.example_ltypes, self.target_ltypes = None, None, None, None    \n",
    "        self.tokenizer = tokenizer\n",
    "        self.lbl_st_len = 0\n",
    "        \n",
    "    def __str__(self):\n",
    "        if(self.child is not None):\n",
    "            return self.child.__str__(spacing=\"\")\n",
    "        else:\n",
    "            return 'empty'\n",
    "    \n",
    "    def get_labelled_set_len(self):\n",
    "        return self.lbl_st_len\n",
    "    \n",
    "    def get_labelled_set(self):\n",
    "        assert self.child is not None\n",
    "        if(self.example_tokens is not None and self.example_ltypes is not None and self.target_tokens is not None and self.target_ltypes is not None):\n",
    "            return self.example_tokens, self.target_tokens, self.example_ltypes,  self.target_ltypes\n",
    "        example_tokens_list, target_tokens_list, example_ltypes_list, target_ltypes_list = [], [], [], []\n",
    "        for child in self.child.get_next_leaf():\n",
    "            example_tokens, target_tokens, example_types, target_ltypes = child.get_labelled_set()\n",
    "            if(example_tokens is not None):\n",
    "                example_tokens_list.append(example_tokens)\n",
    "                target_tokens_list.append(target_tokens)\n",
    "                example_ltypes_list.append(example_types)\n",
    "                target_ltypes_list.append(target_ltypes)\n",
    "        self.example_tokens = torch.cat(example_tokens_list, dim=0)\n",
    "        self.example_ltypes = torch.cat(example_ltypes_list, dim=0)\n",
    "        self.target_tokens = torch.cat(target_tokens_list, dim=0)\n",
    "        self.target_ltypes = torch.cat(target_ltypes_list, dim=0)\n",
    "        self.lbl_st_len = self.example_tokens.shape[0]\n",
    "        return self.example_tokens, self.target_tokens, self.example_ltypes,  self.target_ltypes\n",
    "    \n",
    "    def get_node_labelled_set(self, cur, nxt):\n",
    "        example_tokens, target_tokens, example_ltypes, target_ltypes = [], [], [], []\n",
    "        for i in range(len(cur.tokens)-1):\n",
    "            example_tokens.append(cur.tokens[i])\n",
    "            example_ltypes.append(cur.get_self_leaf_type())\n",
    "            target_tokens.append(cur.tokens[i+1])\n",
    "            target_ltypes.append(cur.get_self_leaf_type())\n",
    "        if(nxt is not None):\n",
    "            example_tokens.append(cur.tokens[len(cur.tokens)-1])\n",
    "            example_ltypes.append(cur.get_self_leaf_type())\n",
    "            target_tokens.append(nxt.tokens[0])\n",
    "            target_ltypes.append(nxt.get_self_leaf_type())\n",
    "        if(target_tokens == []):\n",
    "            return None, None, None, None\n",
    "        return example_tokens, target_tokens, example_ltypes, target_ltypes\n",
    "    \n",
    "    def label_all_nodes(self):\n",
    "        it = iter(self.get_next_leaf())\n",
    "        cur = next(it)\n",
    "        try:\n",
    "            while(True):\n",
    "                nxt = next(it)\n",
    "                example_tokens, target_tokens, example_ltypes, target_ltypes = self.get_node_labelled_set(cur, nxt)\n",
    "                cur.label_tokens(example_tokens, target_tokens)\n",
    "                cur.label_ltypes(example_ltypes, target_ltypes)\n",
    "                cur = nxt\n",
    "        except StopIteration:\n",
    "            example_tokens, target_tokens, example_ltypes, target_ltypes = self.get_node_labelled_set(cur, None)\n",
    "            if(example_tokens is not None and target_tokens is not None):\n",
    "                cur.label_tokens(example_tokens, target_tokens)\n",
    "                cur.label_ltypes(example_ltypes, target_ltypes)\n",
    "    \n",
    "    def generate_text(self):\n",
    "        return self.child.generate_text() if(self.child is not None) else \"\"\n",
    "        \n",
    "    def set_parent(self, parent):\n",
    "        raise Exception('cannot set parent for RootLeaf')\n",
    "        \n",
    "    def replace_child(self, old, new):\n",
    "        if(self.child == old):\n",
    "            self.child = new\n",
    "        else:\n",
    "            raise Exception('does not exist as child')\n",
    "            \n",
    "    def push_token(self, token, wtype):\n",
    "        ltype = self.get_leaf_type_of(wtype)\n",
    "        if(ltype == self.leaf_types_dict[\"ONE_CHILD_LEAF\"]):\n",
    "            child = OneChildLeaf(token, wtype, self.tokenizer, self) \n",
    "        elif(ltype == self.leaf_types_dict[\"TWO_CHILD_LEAF\"]):\n",
    "            child = TwoChildLeaf(token, wtype, self.tokenizer, self)\n",
    "        elif(ltype == self.leaf_types_dict[\"BRIDGE_LEAF\"] or ltype == self.leaf_types_dict[\"INVARIANT\"]):\n",
    "            child = BridgeLeaf(token, wtype, self.tokenizer, self)\n",
    "        else:\n",
    "            raise Exception(f'invalid leaf type passed for token:{token} ltype:{ltype} and wtype:{wtype}')\n",
    "        self.child = child\n",
    "        return child\n",
    "\n",
    "    def get_ordered_children(self):\n",
    "        before, after = [], []\n",
    "        if(self.child is not None): after.append(self.child)\n",
    "        return before, after\n",
    "        \n",
    "    def get_ordered_children_for_model(self):\n",
    "        before, after = [], []\n",
    "        if(self.child is not None): after.append(self.child)\n",
    "        return before, after"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SentenceTree():\n",
    "    def __init__(self, tokenizer, tokens_list):\n",
    "        self.root = RootLeaf(tokenizer)\n",
    "        self.current = self.root\n",
    "        self.model_current = self.root\n",
    "        self.tokenizer, self.tokens_list = tokenizer, tokens_list\n",
    "        self.push_tokens(tokens_list)\n",
    "        self.root.label_all_nodes()\n",
    "        \n",
    "    def __str__(self):\n",
    "        return str(self.root)\n",
    "\n",
    "    def push_token(self, word, wtype):\n",
    "        self.current = self.current.push_token(word, wtype)\n",
    "\n",
    "    def get_labelled_set(self):\n",
    "        return self.root.get_labelled_set()\n",
    "    \n",
    "    def push_tokens(self, token_list):\n",
    "        for token, wtype in token_list:\n",
    "            self.push_token(token, wtype)\n",
    "\n",
    "    def get_beginning(self):\n",
    "        return self.root\n",
    "\n",
    "    def get_current(self):\n",
    "        return sef.current\n",
    "\n",
    "    def generate_text(self):\n",
    "        return self.root.generate_text() \n",
    "    \n",
    "    def get_labelled_set_len(self):\n",
    "        return self.root.get_labelled_set_len()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "class LangTree():\n",
    "    def __init__(self, tokenizer, root_dest_path=Path('./data/lang_tree_data'), bs=32, bptt=70):\n",
    "        root_dest_path = create_dir(make_to_path(root_dest_path))\n",
    "        self.tokenizer,  self.root_dest_path = tokenizer, root_dest_path\n",
    "        self.length = 0\n",
    "        self.batch_dest = root_dest_path / 'batches'\n",
    "        self.batch_dest = create_dir(self.batch_dest)\n",
    "        self.dest_paths = self.prepare_tree_sequenced_tokens()\n",
    "        self.batch_paths = self.batchify(self.dest_paths, self.batch_dest, self.length, bs, bptt)\n",
    "        self.bs, self.bptt = bs, bptt\n",
    "        \n",
    "    def __len__(self):\n",
    "        return self.length\n",
    "    \n",
    "    def get_one_batch_len(self):\n",
    "        return self.one_batch_num_blocks\n",
    "    \n",
    "    def get_num_ltypes(self):\n",
    "        return len(self.tokenizer.get_ltypes_dict())\n",
    "    \n",
    "    def __str__(self):\n",
    "        s = \"\"\n",
    "        for r in self.roots:\n",
    "            s += str(r)\n",
    "            s += \"\\n\\n\"\n",
    "        return s\n",
    "    \n",
    "    def reset_len(self):\n",
    "        self.length = 0\n",
    "        \n",
    "    def get_batch_of_bptt(self, batch_num, seq_num):\n",
    "        with open(self.batch_paths[f\"{batch_num}_{seq_num}\"], \"rb\") as f:\n",
    "            return pickle.load(f)\n",
    "        \n",
    "    def prepare_tree_sequenced_tokens(self):\n",
    "        dest_paths = []\n",
    "        ind = 0\n",
    "        for path, token_list in self.tokenizer.get_next_token_list():\n",
    "            stree = SentenceTree(self.tokenizer, token_list)\n",
    "            tree_tok_sequence = stree.get_labelled_set()\n",
    "            self.length += stree.get_labelled_set_len()\n",
    "            fname = path.stem\n",
    "            dest_path = self.root_dest_path / f\"{fname}_{ind}.pkl\"\n",
    "            dest_paths.append(dest_path)\n",
    "            with open(dest_path, \"wb\") as f:\n",
    "                pickle.dump(tree_tok_sequence, f)\n",
    "            gc.collect()\n",
    "            ind += 1\n",
    "        return dest_paths\n",
    "\n",
    "    def batchify(self, src_paths, dest_root, length, bs, bptt):\n",
    "        def get_next_token_from_paths(src_paths):\n",
    "            for path in src_paths:\n",
    "                with open(path, \"rb\") as f:\n",
    "                    e_toks, t_toks, e_types, t_types = pickle.load(f)\n",
    "                    for i in range(len(e_toks)):\n",
    "                        yield (e_toks[i], t_toks[i], e_types[i], t_types[i])\n",
    "                    \n",
    "        one_batch_len = length // bs\n",
    "        self.one_batch_num_blocks = one_batch_len // bptt\n",
    "        print(length, bs, one_batch_len, bptt)\n",
    "        batch_num = 0\n",
    "        seq_num = 0\n",
    "        batch_paths = {}\n",
    "        next_tok_iter = iter(get_next_token_from_paths(src_paths))\n",
    "        for batch_num in range(bs):\n",
    "            for seq_num in range(one_batch_len // bptt):\n",
    "                dest_path = dest_root / f\"{batch_num}_{seq_num}.pkl\"\n",
    "                batch_paths[f\"{batch_num}_{seq_num}\"] = dest_path\n",
    "                e_tok_list, t_tok_list, e_type_list, t_type_list = [], [], [], []\n",
    "                for s in range(bptt):\n",
    "                    tup = next(next_tok_iter)\n",
    "                    e_tok, t_tok, e_types, t_types = tup\n",
    "                    e_tok_list.append(e_tok)\n",
    "                    t_tok_list.append(t_tok)\n",
    "                    e_type_list.append(e_types)\n",
    "                    t_type_list.append(t_types)\n",
    "                e_toks_t = torch.cat(e_tok_list, dim=0)\n",
    "                t_toks_t = torch.cat(t_tok_list, dim=0)\n",
    "                e_types_t = torch.cat(e_type_list, dim=0)\n",
    "                t_types_t = torch.cat(t_type_list, dim=0)\n",
    "                with open(dest_path, \"wb\") as f:\n",
    "                    pickle.dump((e_toks_t[:,None], t_toks_t[:,None], e_types_t[:,None], t_types_t[:,None]), f)\n",
    "                    print(f\"batch: {batch_num} seq: {seq_num} wrote to \" + str(dest_path))\n",
    "            \n",
    "        return batch_paths\n",
    "                        \n",
    "    def get_next_labelled_set(self, num_tokens):\n",
    "        for path in self.dest_paths:\n",
    "            with open(path, \"rb\") as f:\n",
    "                example_tokens, target_tokens, example_ltypes, target_ltypes = pickle.load(f)\n",
    "            for i in range(0, len(example_tokens), num_tokens):\n",
    "                j = min(i + num_tokens, len(example_tokens))\n",
    "                yield example_tokens[i:j], target_tokens[i:j], example_ltypes[i:j], target_ltypes[i:j]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tk = Tokenizer(text_list=btext_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "it = iter(tk.get_next_token_list())\n",
    "x = next(it)\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "120254 // 32, 3757 // 70"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree = LangTree(tk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test\n",
    "def tree_test(tree):\n",
    "    generated_list = [tree.generate_text()]\n",
    "    for (gen, sen) in zip(generated_list, tree.vocab.get_batched_token_sequences()[0:1]):\n",
    "        s1 = gen.split(' ')[1:]\n",
    "        s2 = [token[0] for token in sen]\n",
    "        assert s1==s2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset, DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_train_valid_indices(tree, val_split=0.1):\n",
    "    tot_size = tree.get_one_batch_len()\n",
    "    val_size = int(0.1*tot_size)\n",
    "    train_size = tot_size - val_size\n",
    "    return [0, train_size], [train_size, tot_size]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "get_train_valid_indices(tree)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TreeDataset(Dataset):\n",
    "    def __init__(self, tree, idx_lower, idx_upper):\n",
    "        self.tree = tree\n",
    "        self.idx_lower, self.idx_upper = idx_lower, idx_upper\n",
    "        self.bs, self.bptt = self.tree.bs, self.tree.bptt\n",
    "    def __len__(self):\n",
    "        return self.idx_upper - self.idx_lower\n",
    "    def __getitem__(self, idx):\n",
    "        e_tok_l, t_tok_l, e_type_l, t_type_l = [], [], [], []\n",
    "        for i in range(tree.bs):\n",
    "            e_toks, t_toks, e_types, t_types = tree.get_batch_of_bptt(i, self.idx_lower + idx)\n",
    "            e_tok_l.append(e_toks)\n",
    "            t_tok_l.append(t_toks)\n",
    "            e_type_l.append(e_types)\n",
    "            t_type_l.append(t_types)\n",
    "        e_tok_t, t_tok_t, e_type_t, t_type_t = torch.stack(e_tok_l), torch.stack(t_tok_l), torch.stack(e_type_l), torch.stack(t_type_l)\n",
    "        return e_tok_t, t_tok_t, e_type_t, t_type_t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TreeDataLoader():\n",
    "    def __init__(self, tree_ds):\n",
    "        self.ds = tree_ds\n",
    "        self.length = len(tree_ds)\n",
    "    def __iter__(self):\n",
    "        for i in range(self.length):\n",
    "            yield self.ds[i]\n",
    "    def __len__(self):\n",
    "        return self.ds.bptt*self.length*self.ds.bs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(\"train:\")\n",
    "# for a,b,c,d in train_dl:\n",
    "#     print(a.shape)\n",
    "# print(\"valid:\")\n",
    "# for a,b,c,d in valid_dl:\n",
    "#     print(a.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train_ind, valid_ind = get_train_valid_indices(tree, 0.1)\n",
    "# train_ds = TreeDataset(tree, *train_ind)\n",
    "# valid_ds = TreeDataset(tree, *valid_ind)\n",
    "# train_dl, valid_dl = TreeDataLoader(train_ds), TreeDataLoader(valid_ds)\n",
    "\n",
    "# it = iter(train_dl)\n",
    "# rnn = nn.RNN(50,100, batch_first=True)\n",
    "# emb = nn.Embedding(len(tree.tokenizer.vocab.itos), 50)\n",
    "# e = emb(a).squeeze()\n",
    "# print(e.shape)\n",
    "# r = rnn(e)[0]\n",
    "# print(r.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SerialTreeModel(nn.Module):\n",
    "    def __init__(self, tree, bs):\n",
    "        super().__init__()\n",
    "        self.emb_tok = nn.Embedding(len(tree.tokenizer.vocab.itos), 50)\n",
    "        self.emb_ltype = nn.Embedding(tree.get_num_ltypes(), 50)\n",
    "        self.rnn = nn.RNN(100,100, batch_first=True)\n",
    "        self.lin_tok = nn.Linear(50, len(tree.tokenizer.vocab.itos))\n",
    "        self.lin_ltype = nn.Linear(50, tree.get_num_ltypes())\n",
    "        self.h = torch.zeros(1, bs, 100)\n",
    "        self.h = self.h.cuda() if torch.cuda.is_available() else self.h\n",
    "        \n",
    "    def forward(self, toks, ltypes):\n",
    "        a1 = self.emb_tok(toks).squeeze()\n",
    "        a2 = self.emb_ltype(ltypes).squeeze()\n",
    "        ip_rnn = torch.cat([a1, a2], dim=-1)\n",
    "        \n",
    "        res, h = self.rnn(ip_rnn, self.h)\n",
    "        self.h = h.detach()\n",
    "        a1, a2 = torch.split(res, 50, dim=-1)\n",
    "        tok_out = self.lin_tok(a1)\n",
    "        ltype_out = self.lin_ltype(a2)\n",
    "        \n",
    "        return tok_out, ltype_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tk = Tokenizer(text_list=btext_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree = LangTree(tk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ind, valid_ind = get_train_valid_indices(tree, 0.1)\n",
    "train_ds = TreeDataset(tree, *train_ind)\n",
    "valid_ds = TreeDataset(tree, *valid_ind)\n",
    "train_dl, valid_dl = TreeDataLoader(train_ds), TreeDataLoader(valid_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SerialTreeModel(tree, 32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_tok, y_tok, x_type, y_type = next(iter(train_dl))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_tok.shape\n",
    "for i in range(x_tok.shape[0]):\n",
    "    print(tree.tokenizer.tokenify(x_tok[i]))\n",
    "    print(tree.tokenizer.tokenify(y_tok[i]), \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(valid_dl), len(train_dl))\n",
    "print(len(tree.tokenizer.vocab.itos))\n",
    "print(model(x_tok, x_type)[0].shape, y_tok.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_func1 = nn.CrossEntropyLoss()\n",
    "y_hat = model(x_tok, x_type)[0]\n",
    "loss =loss_func1(y_hat.squeeze().view(y_hat.shape[0], y_hat.shape[2], y_hat.shape[1]), y_tok.squeeze())\n",
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit(epochs, model, opt, train_dl, valid_dl):\n",
    "    loss_func1 = nn.CrossEntropyLoss()\n",
    "    loss_func2 = nn.CrossEntropyLoss()\n",
    "    for epoch in range(epochs):\n",
    "        model.train()\n",
    "        for x_tok, y_tok, x_ltype, y_ltype in train_dl:\n",
    "            tok_hat, ltype_hat = model(x_tok, x_ltype)\n",
    "            \n",
    "            tok_hat = tok_hat.squeeze().view(tok_hat.shape[0], tok_hat.shape[2], tok_hat.shape[1])\n",
    "            y_tok = y_tok.squeeze()\n",
    "            ltype_hat = ltype_hat.squeeze().view(ltype_hat.shape[0], ltype_hat.shape[2], ltype_hat.shape[1])\n",
    "            y_ltype = y_ltype.squeeze()\n",
    "            \n",
    "            loss1 = loss_func1(tok_hat, y_tok)\n",
    "            loss2 = loss_func2(ltype_hat, y_ltype)\n",
    "            loss = loss1 + loss2\n",
    "            loss.backward()\n",
    "            \n",
    "            opt.step()\n",
    "            opt.zero_grad()\n",
    "\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            tot_loss,tot_acc = 0.,0.\n",
    "            for x_tok_v, y_tok_v, x_ltype_v, y_ltype_v in valid_dl:\n",
    "                tok_hat_v, ltype_hat_v = model(x_tok_v, x_ltype_v)\n",
    "                \n",
    "                tok_hat_v = tok_hat_v.squeeze().view(tok_hat_v.shape[0], tok_hat_v.shape[2], tok_hat_v.shape[1])\n",
    "                y_tok_v = y_tok_v.squeeze()\n",
    "                ltype_hat_v = ltype_hat_v.squeeze().view(ltype_hat_v.shape[0], ltype_hat_v.shape[2], ltype_hat_v.shape[1])\n",
    "                y_ltype_v = y_ltype_v.squeeze()\n",
    "            \n",
    "                loss1_v = loss_func1(tok_hat_v, y_tok_v)\n",
    "                loss2_v = loss_func2(ltype_hat_v, y_ltype_v)\n",
    "                loss_v = loss1_v + loss2_v\n",
    "                tot_loss += loss_v\n",
    "#                 tot_acc  += (accuracy(tok_hat_v, y_tok_v) + accuracy(ltype_hat_v, y_ltype_v))/2\n",
    "        nv = len(valid_dl)\n",
    "        print(epoch, tot_loss/nv)\n",
    "\n",
    "#         print(epoch, tot_loss/nv, tot_acc/nv)\n",
    "    return tot_loss/nv #, tot_acc/nv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fit(10, model, torch.optim.Adam(model.parameters(), 0.01), train_dl, valid_dl)"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}