"""
        """
        super().__init__()

        self.embedding = nn.Linear(elem_emb_len, elem_fea_len)

        self.convs = nn.ModuleList(
            [CGCNNConv(
                elem_fea_len=elem_fea_len,
                nbr_fea_len=nbr_fea_len
            ) for _ in range(n_graph)]
        )

import torch
import torch.nn as nn
import torch.nn.functional as F

from roost.core import BaseModelClass
from roost.segments import MeanPooling, SimpleNetwork, SumPooling

class CrystalGraphConvNet(BaseModelClass):
    """
    Create a crystal graph convolutional neural network for predicting total
    material properties.

    Changes to the code were made to allow for the removal of zero-padding
    and to benefit from the BaseModelClass functionality. The architectural
    choices of the model remain unchanged.
    """

    def __init__(
        self,
        robust,
        n_targets,
        elem_emb_len,
        nbr_fea_len,
        elem_fea_len=64,
        n_graph=4,
        h_fea_len=128,
        n_trunk=1,
        n_hidden=1,
        **kwargs,
    ):
        """
        Initialize CrystalGraphConvNet.

        Parameters
        ----------

        orig_elem_fea_len: int
            Number of atom features in the input.
        nbr_fea_len: int
            Number of bond features.
        elem_fea_len: int
            Number of hidden atom features in the convolutional layers
        n_graph: int
            Number of convolutional layers
        h_fea_len: int
            Number of hidden features after pooling
        n_hidden: int
            Number of hidden layers after pooling
        """
        super().__init__(robust=robust, **kwargs)

        desc_dict = {
            "elem_emb_len": elem_emb_len,
            "nbr_fea_len": nbr_fea_len,
            "elem_fea_len": elem_fea_len,
            "n_graph": n_graph,
        }

        self.node_nn = DescriptorNetwork(**desc_dict)

        self.model_params.update(
            {
                "robust": robust,
                "n_targets": n_targets,
                "h_fea_len": h_fea_len,
                "n_hidden": n_hidden,
            }
        )

        self.model_params.update(desc_dict)

        self.pooling = MeanPooling()

        # define an output neural network
        if self.robust:
            n_targets = [2 * n for n in n_targets]

        out_hidden = [h_fea_len] * n_hidden
        trunk_hidden = [h_fea_len] * n_trunk
        self.trunk_nn = SimpleNetwork(elem_fea_len, h_fea_len, trunk_hidden)

        self.output_nns = nn.ModuleList(
            SimpleNetwork(h_fea_len, n, out_hidden) for n in n_targets
        )

    def forward(self, atom_fea, nbr_fea, self_idx, nbr_idx, crystal_atom_idx):
        """
        Forward pass

        N: Total number of atoms in the batch
        M: Max number of neighbors
        N0: Total number of crystals in the batch

        Parameters
        ----------

        atom_fea: Variable(torch.Tensor) shape (N, orig_elem_fea_len)
            Atom features from atom type
        nbr_fea: Variable(torch.Tensor) shape (N, M, nbr_fea_len)
            Bond features of each atom's M neighbors
        nbr_fea_idx: torch.LongTensor shape (N, M)
            Indices of M neighbors of each atom
        crystal_atom_idx: list of torch.LongTensor of length N0
            Mapping from the crystal idx to atom idx

        Returns
        -------

        prediction: nn.Variable shape (N, )
            Atom hidden features after convolution

        """
        atom_fea = self.node_nn(
            atom_fea, nbr_fea, self_idx, nbr_idx
        )

        crys_fea = self.pooling(atom_fea, crystal_atom_idx)

        # NOTE required to match the reference implementation
        crys_fea = nn.functional.softplus(crys_fea)

        crys_fea = F.relu(self.trunk_nn(crys_fea))

        # apply neural network to map from learned features to target
        return (output_nn(crys_fea) for output_nn in self.output_nns)

class DescriptorNetwork(nn.Module):
    """
    The Descriptor Network is the message passing section of the
    CrystalGraphConvNet Model.
    """

    def __init__(
        self, elem_emb_len, nbr_fea_len, elem_fea_len=64, n_graph=4,
    ):

    def forward(self, atom_fea, nbr_fea, self_fea_idx, nbr_fea_idx):
        """
        Forward pass

        N: Total number of atoms in the batch
        M: Max number of neighbors
        N0: Total number of crystals in the batch

        Parameters
        ----------

        atom_fea: Variable(torch.Tensor) shape (N, orig_elem_fea_len)
            Atom features from atom type
        nbr_fea: Variable(torch.Tensor) shape (N, M, nbr_fea_len)
            Bond features of each atom's M neighbors
        nbr_fea_idx: torch.LongTensor shape (N, M)
            Indices of M neighbors of each atom
        crystal_atom_idx: list of torch.LongTensor of length N0
            Mapping from the crystal idx to atom idx

        Returns
        -------

        prediction: nn.Variable shape (N, )
            Atom hidden features after convolution

        """
        atom_fea = self.embedding(atom_fea)

        for conv_func in self.convs:
            atom_fea = conv_func(atom_fea, nbr_fea, self_fea_idx, nbr_fea_idx)

        return atom_fea

class CGCNNConv(nn.Module):
    """
    Convolutional operation on graphs
    """

    def __init__(self, elem_fea_len, nbr_fea_len):
        """
        Initialize CGCNNConv.

        Parameters
        ----------

        elem_fea_len: int
                Number of atom hidden features.
        nbr_fea_len: int
                Number of bond features.
        """
        super().__init__()
        self.elem_fea_len = elem_fea_len
        self.nbr_fea_len = nbr_fea_len
        self.fc_full = nn.Linear(
            2 * self.elem_fea_len + self.nbr_fea_len, 2 * self.elem_fea_len
        )
        self.sigmoid = nn.Sigmoid()
        self.softplus1 = nn.Softplus()
        self.bn1 = nn.BatchNorm1d(2 * self.elem_fea_len)
        self.bn2 = nn.BatchNorm1d(self.elem_fea_len)
        self.softplus2 = nn.Softplus()
        self.pooling = SumPooling()

    def forward(self, atom_in_fea, nbr_fea, self_fea_idx, nbr_fea_idx):
        """
        Forward pass

        N: Total number of atoms in the batch
        M: Max number of neighbors

        Parameters
        ----------

        atom_in_fea: Variable(torch.Tensor) shape (N, elem_fea_len)
            Atom hidden features before convolution
        nbr_fea: Variable(torch.Tensor) shape (N, M, nbr_fea_len)
            Bond features of each atom's M neighbors
        nbr_fea_idx: torch.LongTensor shape (N, M)
            Indices of M neighbors of each atom

        Returns
        -------

        atom_out_fea: nn.Variable shape (N, elem_fea_len)
            Atom hidden features after convolution

        """
        # convolution
        atom_nbr_fea = atom_in_fea[nbr_fea_idx, :]
        atom_self_fea = atom_in_fea[self_fea_idx, :]

        total_fea = torch.cat([atom_self_fea, atom_nbr_fea, nbr_fea], dim=1)

        total_fea = self.fc_full(total_fea)
        total_fea = self.bn1(total_fea)

        filter_fea, core_fea = total_fea.chunk(2, dim=1)
        filter_fea = self.sigmoid(filter_fea)
        core_fea = self.softplus1(core_fea)

        # take the elementwise product of the filter and core
        nbr_msg = filter_fea * core_fea
        nbr_sumed = self.pooling(nbr_msg, self_fea_idx)

        nbr_sumed = self.bn2(nbr_sumed)
        out = self.softplus2(atom_in_fea + nbr_sumed)

        return out


        super().__init__()

        self.embedding = nn.Linear(elem_emb_len, elem_fea_len)

        self.convs = nn.ModuleList(
            [CGCNNConv(
                elem_fea_len=elem_fea_len,
                nbr_fea_len=nbr_fea_len
            ) for _ in range(n_graph)]
        )

import torch
import torch.nn as nn
import torch.nn.functional as F

from roost.core import BaseModelClass
from roost.segments import MeanPooling, SimpleNetwork, SumPooling

class CrystalGraphConvNet(BaseModelClass):
    """
    Create a crystal graph convolutional neural network for predicting total
    material properties.

    Changes to the code were made to allow for the removal of zero-padding
    and to benefit from the BaseModelClass functionality. The architectural
    choices of the model remain unchanged.
    """

    def __init__(
        self,
        robust,
        n_targets,
        elem_emb_len,
        nbr_fea_len,
        elem_fea_len=64,
        n_graph=4,
        h_fea_len=128,
        n_trunk=1,
        n_hidden=1,
        **kwargs,
    ):
        """
        Initialize CrystalGraphConvNet.

        Parameters
        ----------

        orig_elem_fea_len: int
            Number of atom features in the input.
        nbr_fea_len: int
            Number of bond features.
        elem_fea_len: int
            Number of hidden atom features in the convolutional layers
        n_graph: int
            Number of convolutional layers
        h_fea_len: int
            Number of hidden features after pooling
        n_hidden: int
            Number of hidden layers after pooling
        """
        super().__init__(robust=robust, **kwargs)

        desc_dict = {
            "elem_emb_len": elem_emb_len,
            "nbr_fea_len": nbr_fea_len,
            "elem_fea_len": elem_fea_len,
            "n_graph": n_graph,
        }

        self.node_nn = DescriptorNetwork(**desc_dict)

        self.model_params.update(
            {
                "robust": robust,
                "n_targets": n_targets,
                "h_fea_len": h_fea_len,
                "n_hidden": n_hidden,
            }
        )

        self.model_params.update(desc_dict)

        self.pooling = MeanPooling()

        # define an output neural network
        if self.robust:
            n_targets = [2 * n for n in n_targets]

        out_hidden = [h_fea_len] * n_hidden
        trunk_hidden = [h_fea_len] * n_trunk
        self.trunk_nn = SimpleNetwork(elem_fea_len, h_fea_len, trunk_hidden)

        self.output_nns = nn.ModuleList(
            SimpleNetwork(h_fea_len, n, out_hidden) for n in n_targets
        )

    def forward(self, atom_fea, nbr_fea, self_idx, nbr_idx, crystal_atom_idx):
        """
        Forward pass

        N: Total number of atoms in the batch
        M: Max number of neighbors
        N0: Total number of crystals in the batch

        Parameters
        ----------

        atom_fea: Variable(torch.Tensor) shape (N, orig_elem_fea_len)
            Atom features from atom type
        nbr_fea: Variable(torch.Tensor) shape (N, M, nbr_fea_len)
            Bond features of each atom's M neighbors
        nbr_fea_idx: torch.LongTensor shape (N, M)
            Indices of M neighbors of each atom
        crystal_atom_idx: list of torch.LongTensor of length N0
            Mapping from the crystal idx to atom idx

        Returns
        -------

        prediction: nn.Variable shape (N, )
            Atom hidden features after convolution

        """
        atom_fea = self.node_nn(
            atom_fea, nbr_fea, self_idx, nbr_idx
        )

        crys_fea = self.pooling(atom_fea, crystal_atom_idx)

        # NOTE required to match the reference implementation
        crys_fea = nn.functional.softplus(crys_fea)

        crys_fea = F.relu(self.trunk_nn(crys_fea))

        # apply neural network to map from learned features to target
        return (output_nn(crys_fea) for output_nn in self.output_nns)

class DescriptorNetwork(nn.Module):
    """
    The Descriptor Network is the message passing section of the
    CrystalGraphConvNet Model.
    """

    def __init__(
        self, elem_emb_len, nbr_fea_len, elem_fea_len=64, n_graph=4,
    ):
        """
        """


if request.method in permissions.SAFE_METHODS:
            return True
        return request.user and request.user.is_staff

from rest_framework import permissions

class IsAdminUserOrReadOnly(permissions.BasePermission):
    def has_permission(self, request, view):


import asyncio

import uvloop
from sanic import Sanic

from momo.settings import Config

blueprints = []

def create_app(register_bp=True, test=False):
    app = Sanic(__name__)
    if test:
        app.config['TESTING'] = True
    app.config.from_object(Config)
    register_blueprints(app)
    return app

def register_extensions(app):

def register_blueprints(app):
    from momo.views.hello import blueprint as hello_bp
    from momo.views.mweixin import blueprint as wx_bp
    app.register_blueprint(hello_bp)
    app.register_blueprint(wx_bp)

def register_jinja_funcs(app):
    return app

app = create_app()
asyncio.set_event_loop(uvloop.new_event_loop())
server = app.create_server(host="0.0.0.0", port=8888, debug=True)
loop = asyncio.get_event_loop()
task = asyncio.ensure_future(server)
try:
    loop.run_forever()
except:
    loop.stop()


try:
    loop.run_forever()
except:
    loop.stop()

import asyncio

import uvloop
from sanic import Sanic

from momo.settings import Config

blueprints = []

def create_app(register_bp=True, test=False):
    app = Sanic(__name__)
    if test:
        app.config['TESTING'] = True
    app.config.from_object(Config)
    register_blueprints(app)
    return app

def register_extensions(app):
    pass

def register_blueprints(app):
    from momo.views.hello import blueprint as hello_bp
    from momo.views.mweixin import blueprint as wx_bp
    app.register_blueprint(hello_bp)
    app.register_blueprint(wx_bp)

def register_jinja_funcs(app):
    return app

app = create_app()
asyncio.set_event_loop(uvloop.new_event_loop())
server = app.create_server(host="0.0.0.0", port=8888, debug=True)
loop = asyncio.get_event_loop()
task = asyncio.ensure_future(server)


self._links = [x.replace('\n', '') for x in self.framework.reglib(self._pages).search(r'\n(https?://.*)')]
		self._links.extend([x.replace('class="link">\n', '') for x in self.framework.reglib(self._pages).search(r'class="link">\n(https?://.*)')])
		return self._links

"""
OWASP Maryam!

"""

class main:

	def __init__(self, q, limit):
		""" use onionlandsearchengine.com

			q 		  : Query for search
			limit	  : Number of pages
		"""
		self.framework = main.framework
		self.q = self.framework.urlib(q).quote
		self.limit = 20 if limit > 20 else limit
		self._pages = ''
		self._links = []
		self.onionlandsearchengine = 'onionlandsearchengine.com'

	def run_crawl(self):
		urls = [f"https://{self.onionlandsearchengine}/search?q={self.q}&page={i}" for i in range(1, self.limit)]
		max_attempt = len(urls)
		plen = 0
		for url in range(max_attempt):
			self.framework.debug(f'[ONIONLAND] Searching in {url} page...')
			try:
				req = self.framework.request(url=urls[url])
			except:
				self.framework.print_exception()
				max_attempt -= 1
				if max_attempt == 0:
					self.framework.error('Onionland is missed!', 'util/onionland', 'run_crawl')
					break
			else:
				self._pages += req.text
				if plen == 0:
					plen = len(self.framework.reglib(self._pages).search('class="page"'))
				if plen-1 == url:
					return

	@property
	def pages(self):
		return self._pages

	@property
	def links(self):


		self.framework = main.framework
		self.q = self.framework.urlib(q).quote
		self.limit = 20 if limit > 20 else limit
		self._pages = ''
		self._links = []
		self.onionlandsearchengine = 'onionlandsearchengine.com'

"""
OWASP Maryam!

"""

class main:

	def __init__(self, q, limit):
		""" use onionlandsearchengine.com

			q 		  : Query for search
			limit	  : Number of pages
		"""

def run_crawl(self):
		urls = [f"https://{self.onionlandsearchengine}/search?q={self.q}&page={i}" for i in range(1, self.limit)]
		max_attempt = len(urls)
		plen = 0
		for url in range(max_attempt):
			self.framework.debug(f'[ONIONLAND] Searching in {url} page...')
			try:
				req = self.framework.request(url=urls[url])
			except:
				self.framework.print_exception()
				max_attempt -= 1
				if max_attempt == 0:
					self.framework.error('Onionland is missed!', 'util/onionland', 'run_crawl')
					break
			else:
				self._pages += req.text
				if plen == 0:
					plen = len(self.framework.reglib(self._pages).search('class="page"'))
				if plen-1 == url:
					return

	@property
	def pages(self):
		return self._pages

	@property
	def links(self):
		self._links = [x.replace('\n', '') for x in self.framework.reglib(self._pages).search(r'\n(https?://.*)')]
		self._links.extend([x.replace('class="link">\n', '') for x in self.framework.reglib(self._pages).search(r'class="link">\n(https?://.*)')])
		return self._links


try:
				req = self.framework.request(url=urls[url])
			except:
				self.framework.print_exception()
				max_attempt -= 1
				if max_attempt == 0:
					self.framework.error('Onionland is missed!', 'util/onionland', 'run_crawl')
					break
			else:
				self._pages += req.text
				if plen == 0:
					plen = len(self.framework.reglib(self._pages).search('class="page"'))
				if plen-1 == url:
					return

"""
OWASP Maryam!

"""

class main:

	def __init__(self, q, limit):
		""" use onionlandsearchengine.com

			q 		  : Query for search
			limit	  : Number of pages
		"""
		self.framework = main.framework
		self.q = self.framework.urlib(q).quote
		self.limit = 20 if limit > 20 else limit
		self._pages = ''
		self._links = []
		self.onionlandsearchengine = 'onionlandsearchengine.com'

	def run_crawl(self):
		urls = [f"https://{self.onionlandsearchengine}/search?q={self.q}&page={i}" for i in range(1, self.limit)]
		max_attempt = len(urls)
		plen = 0
		for url in range(max_attempt):
			self.framework.debug(f'[ONIONLAND] Searching in {url} page...')

@property
	def pages(self):
		return self._pages

	@property
	def links(self):
		self._links = [x.replace('\n', '') for x in self.framework.reglib(self._pages).search(r'\n(https?://.*)')]
		self._links.extend([x.replace('class="link">\n', '') for x in self.framework.reglib(self._pages).search(r'class="link">\n(https?://.*)')])
		return self._links


from .. import Provider as AddressProvider

class Provider(AddressProvider):

    city_formats = ("{{city_name}}",)

    street_name_formats = (
        "{{first_name}} {{last_name}}",
        "{{last_name}}",
    )

    street_address_formats = ("{{building_number}} {{street_name}}",)

    address_formats = (
        "{{street_address}}\n{{city}} {{postcode}}",
        "{{street_address}}\n{{city}}-{{postcode}}",
    )

    building_number_formats = (
        "####",
        "###",
        "##",
        "#",
        "#/#",
        "##/##",
        "##/###",
        "##/####",
    )

    postcode_formats = ("######",)

    cities = (
        "சென்னை",
        "கோயம்புத்தூர்",
        "மதுரை",
        "திருச்சிராப்பள்ளி",
        "திருப்பூர்",
        "சேலம்",
        "ஈரோடு",
        "திருநெல்வேலி",
        "வேலூர்",
        "தூத்துக்குடி",
        "திண்டுக்கல்",
        "தஞ்சாவூர்",
        "இராணிப்பேட்டை",
        "சிவகாசி",
        "கரூர் (கரூர் மாவட்டம்)",
        "உதகமண்டலம்",
        "ஓசூர்",
        "நாகர்கோவில்",
        "காஞ்சிபுரம்",
        "குமாரபாளையம்",
        "காரைக்குடி",
        "நெய்வேலி",
        "கடலூர்",
        "கும்பகோணம்",
        "திருவண்ணாமலை",
        "பொள்ளாச்சி",
        "இராஜபாளையம், விருதுநகர் மாவட்டம்",
        "குடியாத்தம்",
        "புதுக்கோட்டை",
        "வாணியம்பாடி",
        "ஆம்பூர்",
        "நாகப்பட்டினம்",
        "மும்பை பெருநகர்",
        "தில்லி",
        "கொல்கத்தா பெருநகர்",
        "சென்னை பெருநகர்",
        "பெங்களூரு",
        "ஐதராபாத்",
        "புனே",
        "அகமதாபாத்",
        "கான்பூர்",
        "சூரத்",
        "ஜெய்ப்பூர்",
        "லக்னோ",
        "பாட்னா",
        "நாக்பூர்",
        "இந்தோர்",
        "மீரட்",
        "நாசிக்",
        "போபால்",
        "லூதியானா",
        "ஆக்ரா",
        "வதோதரா",
        "புவனேசுவர்",
        "கோயம்புத்தூர்",
        "ராஜ்கோட்",
        "கொச்சி",
        "விசாகப்பட்டினம்",
        "வாரணாசி",
        "மதுரை",
        "ஆசன்சோல்",
        "அலகாபாத்",
        "மைசூர்",
        "ஜபல்பூர்",
        "ஜம்சேத்பூர்",
        "அவுரங்கபாத்",
        "அம்ரித்சர்",
        "தன்பாத்",
        "விஜயவாடா",
        "சோலாப்பூர்",
        "பிலாய்",
        "ஸ்ரீநகர்",
        "ராஞ்சி",
        "திருவனந்தபுரம்",
        "சண்டிகர்",
        "குவஹாத்தி",
        "கோழிக்கோடு",
        "ஜோத்பூர்",
        "குவாலியர்",
        "ஜலந்தர்",
        "திருச்சிராப்பள்ளி",
        "பரேலி",
        "ஹுப்ளி-தர்வாத்",
        "அலிகார்",
        "கோட்டா",
        "மொரதாபாத்",
        "ராய்ப்பூர்",
        "தேராதூன்",
        "கோரக்பூர்",
        "ஜம்மு",
        "அமராவதி",
        "வாரங்கல்",
        "ஜாம்நகர்",
        "பிகானேர்",
        "சாங்கலி",
        "திருப்பூர்",
        "பாவ்நகர்",
        "மங்களூர்",
        "அஜ்மீர்",
        "பொகாரோ",
        "பெல்காம்",
        "புதுச்சேரி",
        "சிலிகுரி",
        "கண்ணூர்",
        "கோலாப்பூர்",
        "நான்தேட்",
        "ரூர்கேலா",
        "துர்காபூர்",
        "குல்பர்கா",
        "குண்டூர்",
        "ஜான்சி",
        "சகாரன்பூர்",
        "கரக்பூர்",
        "கயா",
        "ஜல்கான்",
        "மதுரா",
        "கொல்லம்",
        "கோர்பா",
        "பிரோசாபாத்",
        "திருநெல்வேலி",
        "உஜ்ஜைன்",
        "அகமத்நகர்",
        "நெல்லூர்",
        "ராமகுண்டம்",
        "ராஜமுந்திரி",
        "மாலேகான்",
        "உதயப்பூர்",
        "அகோலா",
        "தாவண்கரே",
        "வேலூர்",
        "திருவண்ணாமலை",
        "காஜுவாகா",
    )

    states = (
        "ஆந்திரப் பிரதேசம்",
        "அருணாச்சலப் பிரதேசம்",
        "அசாம்",
        "பீகார்",
        "சத்தீஸ்கர்",
        "கோவா",
        "குஜராத்",
        "அரியானா",
        "இமாச்சலப் பிரதேசம்",
        "சம்மு காசுமீர்",
        "ஜார்கண்ட்",
        "கர்நாடகா",
        "கேரளா",
        "மத்தியப் பிரதேசம்",
        "மகாராஷ்டிரா",
        "மணிப்பூர்",
        "மேகாலயா",
        "மிசோரம்",
        "நாகலாந்து",
        "ஒரிசா",
        "பஞ்சாப்",
        "ராஜஸ்தான்",
        "சிக்கிம்",
        "தமிழ்நாடு",
        "தெலுங்கானா",
        "திரிபுரா",
        "உத்தரப்பிரதேசம்",
        "உத்தரகண்ட்",
        "மேற்கு வங்கம்",
    )

    countries = (
        "ஆப்கானித்தான்",
        "அல்பேனியா",
        "அல்ஜீரியா",
        "அந்தோரா",
        "அங்கோலா",
        "அன்டிகுவா பர்புடா",
        "அர்கெந்தீனா",
        "ஆர்மீனியா",
        "ஆத்திரேலியா",
        "ஆஸ்திரியா",
        "அசர்பைஜான்",
        "பஹமாஸ்",
        "பகுரைன்",
        "வங்காளதேசம்",
        "பார்படோசு",
        "பெலருஸ்",
        "பெல்ஜியம்",
        "பெலீசு",
        "பெனின்",
        "பூட்டான்",
        "பொலிவியா",
        "பொசுனியா எர்செகோவினா",
        "போட்சுவானா",
        "பிரேசில்",
        "புரூணை",
        "பல்கேரியா",
        "புர்க்கினா பாசோ",
        "புருண்டி",
        "கம்போடியா",
        "கமரூன்",
        "கனடா",
        "கேப் வர்டி",
        "மத்திய ஆப்பிரிக்கக் குடியரசு",
        "சாட்",
        "சிலி",
        "சீனா",
        "கொலம்பியா",
        "கொமொரோசு",
        "காங்கோ மக்களாட்சிக் குடியரசு",
        "காங்கோ மக்களாட்சிக் குடியரசு",
        "கோஸ்ட்டா ரிக்கா",
        "ஐவரி கோஸ்ட்",
        "குரோவாசியா",
        "கியூபா",
        "சைப்பிரசு",
        "செக் குடியரசு",
        "டென்மார்க்",
        "சீபூத்தீ",
        "டொமினிக்கா",
        "டொமினிக்கன் குடியரசு",
        "எக்குவடோர்",
        "எகிப்து",
        "எல் சல்வடோர",
        "எக்குவடோரியல் கினி",
        "எரித்திரியா",
        "எசுத்தோனியா",
        "எதியோப்பியா",
        "பிஜி",
        "பின்லாந்து",
        "பிரான்சு",
        "காபொன்",
        "கம்பியா",
        "சியார்சியா",
        "செருமனி",
        "கானா",
        "கிரேக்க நாடு",
        "கிரெனடா",
        "குவாத்தமாலா",
        "கினியா",
        "கினி-பிசாவு",
        "கயானா",
        "எயிட்டி",
        "ஒண்டுராசு",
        "அங்கேரி",
        "ஐசுலாந்து",
        "இந்தியா",
        "இந்தோனேசியா",
        "ஈரான்",
        "ஈராக்",
        "அயர்லாந்து",
        "இசுரேல்",
        "இத்தாலி",
        "ஜமேக்கா",
        "சப்பான்",
        "யோர்தான்",
        "கசக்கஸ்தான்",
        "கென்யா",
        "கிரிபட்டி",
        "வட கொரியா",
        "தென் கொரியா",
        "குவைத்",
        "கிர்கிசுத்தான்",
        "லாவோஸ்",
        "லாத்வியா",
        "லெபனான்",
        "லெசோத்தோ",
        "லைபீரியா",
        "லிபியா",
        "லீக்கின்ஸ்டைன்",
        "லித்துவேனியா",
        "லக்சம்பர்க்",
        "மாக்கடோனியக் குடியரசு",
        "மடகாசுகர்",
        "மலாவி",
        "மலேசியா",
        "மாலைத்தீவுகள்",
        "மாலி",
        "மால்ட்டா",
        "மார்சல் தீவுகள்",
        "மூரித்தானியா",
        "மொரிசியசு",
        "மெக்சிக்கோ",
        "மைக்குரோனீசியக் கூட்டு நாடுகள்",
        "மல்தோவா",
        "மொனாகோ",
        "மங்கோலியா",
        "மொண்டெனேகுரோ",
        "மொரோக்கோ",
        "மொசாம்பிக்",
        "மியான்மர்",
        "நமீபியா",
        "நவூரு",
        "நேபாளம்",
        "நெதர்லாந்து",
        "நியூசிலாந்து",
        "நிக்கராகுவா",
        "நைஜர்",
        "நைஜீரியா",
        "நோர்வே",
        "ஓமான்",
        "பாக்கித்தான்",
        "பலாவு",
        "பலத்தீன்",
        "பனாமா",
        "பப்புவா நியூ கினி",
        "பரகுவை",
        "பெரு",
        "பிலிப்பீன்சு",
        "போலந்து",
        "போர்த்துகல்",
        "கட்டார்",
        "உருமேனியா",
        "உருசியா",
        "ருவாண்டா",
        "செயிண்ட். கிட்ஸ் நெவிஸ்",
        "செயிண்ட். லூசியா",
        "செயின்ட் வின்செண்டு மற்றும் கிரெனடீன்கள்",
        "சமோவா",
        "சான் மரீனோ",
        "சாவோ தொமே மற்றும் பிரின்சிப்பி",
        "சவூதி அரேபியா",
        "செனிகல்",
        "செர்பியா",
        "சீசெல்சு",
        "சியேரா லியோனி",
        "சிங்கப்பூர்",
        "சிலவாக்கியா",
        "சுலோவீனியா",
        "சொலமன் தீவுகள்",
        "சோமாலியா",
        "தென்னாப்பிரிக்கா",
        "தெற்கு சூடான்",
        "எசுப்பானியா",
        "இலங்கை",
        "சூடான்",
        "சுரிநாம்",
        "சுவாசிலாந்து",
        "சுவீடன்",
        "சுவிட்சர்லாந்து",
        "சிரியா",
        "சீனக் குடியரசு",
        "தாஜிக்ஸ்தான்",
        "தன்சானியா",
        "தாய்லாந்து",
        "கிழக்குத் திமோர்",
        "டோகோ",
        "தொங்கா",
        "டிரினிடாட் மற்றும் டொபாகோ",
        "தூனிசியா",
        "துருக்கி",
        "துருக்மெனிஸ்தான்",
        "துவாலு",
        "உகாண்டா",
        "உக்ரைன்",
        "ஐக்கிய அரபு அமீரகம்",
        "ஐக்கிய இராச்சியம்",
        "ஐக்கிய அமெரிக்கா",
        "உருகுவை",
        "உஸ்பெகிஸ்தான்",
        "வனுவாட்டு",
        "வெனிசுவேலா",
        "வியட்நாம்",
        "மேற்கு சகாரா (Sahrawi)",
        "யேமன்",
        "சாம்பியா",
        "சிம்பாப்வே",
        "அங்கியுலா (UK)",
        "அரூபா (Netherlands)",
        "பெர்முடா (UK)",
        "கேமன் தீவுகள் (UK)",
        "குயெர்ன்சி (கால்வாய் தீவுகள், UK)",
        "யேர்சி (கால்வாய் தீவுகள், UK)",
        "குக் தீவுகள் (New Zealand)",
        "குராசோ (Netherlands)",
        "போக்லாந்து தீவுகள்/Malvinas",
        "பரோயே தீவுகள் (Denmark)",
        "கிப்ரல்டார் (UK)",
        "கிறீன்லாந்து (Denmark)",
        "குவாதலூப்பு (France)",
        "குவாம் (USA)",
        "பிரெஞ்சு கயானா",
        "ஆங்காங்",
        "மாண் தீவு (UK)",
        "கொசோவோ",
        "மக்காவு",
        "மர்தினிக்கு (France)",
        "மயோட்டே (France)",
        "மொன்செராட்",
    )

    def city_name(self) -> str:
        return self.random_element(self.cities)

    def administrative_unit(self) -> str:


while 1:
                    chunk = stream.Read(4096)
                    if not chunk:
                        break
                    bytes += len(chunk)

import pythoncom
import win32con

formats = """CF_TEXT CF_BITMAP CF_METAFILEPICT CF_SYLK CF_DIF CF_TIFF
            CF_OEMTEXT CF_DIB CF_PALETTE CF_PENDATA CF_RIFF CF_WAVE
            CF_UNICODETEXT CF_ENHMETAFILE CF_HDROP CF_LOCALE CF_MAX
            CF_OWNERDISPLAY CF_DSPTEXT CF_DSPBITMAP CF_DSPMETAFILEPICT
            CF_DSPENHMETAFILE""".split()
format_name_map = {}
for f in formats:
    val = getattr(win32con, f)
    format_name_map[val]=f

tymeds = [attr for attr in pythoncom.__dict__.iterkeys() if attr.startswith("TYMED_")]

def DumpClipboard():
    do = pythoncom.OleGetClipboard()
    print "Dumping all clipboard formats..."
    for fe in do.EnumFormatEtc():
        fmt, td, aspect, index, tymed = fe
        tymeds_this = [getattr(pythoncom, t) for t in tymeds if tymed & getattr(pythoncom, t)]
        print "Clipboard format", format_name_map.get(fmt,str(fmt))
        for t_this in tymeds_this:
            # As we are enumerating there should be no need to call
            # QueryGetData, but we do anyway!
            fetc_query = fmt, td, aspect, index, t_this
            try:
                do.QueryGetData(fetc_query)
            except pythoncom.com_error:
                print "Eeek - QGD indicated failure for tymed", t_this
            # now actually get it.
            try:
                medium = do.GetData(fetc_query)
            except pythoncom.com_error, exc:
                print "Failed to get the clipboard data:", exc
                continue
            if medium.tymed==pythoncom.TYMED_GDI:
                data = "GDI handle %d" % medium.data
            elif medium.tymed==pythoncom.TYMED_MFPICT:
                data = "METAFILE handle %d" % medium.data
            elif medium.tymed==pythoncom.TYMED_ENHMF:
                data = "ENHMETAFILE handle %d" % medium.data
            elif medium.tymed==pythoncom.TYMED_HGLOBAL:
                data = "%d bytes via HGLOBAL" % len(medium.data)
            elif medium.tymed==pythoncom.TYMED_FILE:
                data = "filename '%s'" % data
            elif medium.tymed==pythoncom.TYMED_ISTREAM:
                stream = medium.data
                stream.Seek(0,0)
                bytes = 0

                data = "%d bytes via IStream" % bytes
            elif medium.tymed==pythoncom.TYMED_ISTORAGE:
                data = "a IStorage"
            else:
                data = "*** unknown tymed!"
            print " -> got", data
    do = None

if __name__=='__main__':
    DumpClipboard()
    if pythoncom._GetInterfaceCount()+pythoncom._GetGatewayCount():
        print "XXX - Leaving with %d/%d COM objects alive" % \
              (pythoncom._GetInterfaceCount(), pythoncom._GetGatewayCount())


try:
                medium = do.GetData(fetc_query)
            except pythoncom.com_error, exc:
                print "Failed to get the clipboard data:", exc
                continue

import pythoncom
import win32con

formats = """CF_TEXT CF_BITMAP CF_METAFILEPICT CF_SYLK CF_DIF CF_TIFF
            CF_OEMTEXT CF_DIB CF_PALETTE CF_PENDATA CF_RIFF CF_WAVE
            CF_UNICODETEXT CF_ENHMETAFILE CF_HDROP CF_LOCALE CF_MAX
            CF_OWNERDISPLAY CF_DSPTEXT CF_DSPBITMAP CF_DSPMETAFILEPICT
            CF_DSPENHMETAFILE""".split()
format_name_map = {}
for f in formats:
    val = getattr(win32con, f)
    format_name_map[val]=f

tymeds = [attr for attr in pythoncom.__dict__.iterkeys() if attr.startswith("TYMED_")]

def DumpClipboard():
    do = pythoncom.OleGetClipboard()
    print "Dumping all clipboard formats..."
    for fe in do.EnumFormatEtc():
        fmt, td, aspect, index, tymed = fe
        tymeds_this = [getattr(pythoncom, t) for t in tymeds if tymed & getattr(pythoncom, t)]
        print "Clipboard format", format_name_map.get(fmt,str(fmt))
        for t_this in tymeds_this:
            # As we are enumerating there should be no need to call
            # QueryGetData, but we do anyway!
            fetc_query = fmt, td, aspect, index, t_this
            try:
                do.QueryGetData(fetc_query)
            except pythoncom.com_error:
                print "Eeek - QGD indicated failure for tymed", t_this
            # now actually get it.

            if medium.tymed==pythoncom.TYMED_GDI:
                data = "GDI handle %d" % medium.data
            elif medium.tymed==pythoncom.TYMED_MFPICT:
                data = "METAFILE handle %d" % medium.data
            elif medium.tymed==pythoncom.TYMED_ENHMF:
                data = "ENHMETAFILE handle %d" % medium.data
            elif medium.tymed==pythoncom.TYMED_HGLOBAL:
                data = "%d bytes via HGLOBAL" % len(medium.data)
            elif medium.tymed==pythoncom.TYMED_FILE:
                data = "filename '%s'" % data
            elif medium.tymed==pythoncom.TYMED_ISTREAM:
                stream = medium.data
                stream.Seek(0,0)
                bytes = 0
                while 1:
                    chunk = stream.Read(4096)
                    if not chunk:
                        break
                    bytes += len(chunk)
                data = "%d bytes via IStream" % bytes
            elif medium.tymed==pythoncom.TYMED_ISTORAGE:
                data = "a IStorage"
            else:
                data = "*** unknown tymed!"
            print " -> got", data
    do = None

if __name__=='__main__':
    DumpClipboard()
    if pythoncom._GetInterfaceCount()+pythoncom._GetGatewayCount():
        print "XXX - Leaving with %d/%d COM objects alive" % \
              (pythoncom._GetInterfaceCount(), pythoncom._GetGatewayCount())


"""Return an example's size as a float or tuple. This value is used when
        filtering a dataset with ``--max-positions``."""
        return self.src_sizes[index]

import numpy as np
from ncc.data.tools import data_utils
from ncc.data.ncc_dataset import NccDataset
import jsbeautifier
from dataset.augmented_javascript.utils.util import normalize_program
import torch

TYPED_MARKER_START = "__LS__"
TYPED_MARKER_MID = "__LM__"
TYPED_MARKER_END = "__LE__"

def collate(samples, pad_idx, no_type_id):
    if len(samples) == 0:
        return {}

    def merge(key):
        return data_utils.collate_tokens(
            [s[key] for s in samples],
            pad_idx,
        )
    B = len(samples)
    X = merge('subword_ids')
    Y = [s['label_segments'] for s in samples]

    L = X.size(1)

    # Create tensor of sequence lengths, [B]
    lengths = torch.LongTensor([s['subword_ids'].numel() for s in samples])
    # Make masks for each label interval
    labels = torch.zeros(B, L, dtype=torch.long)
    labels.fill_(no_type_id)
    output_attn = torch.zeros(B, L, L, dtype=torch.float)
    for i, y in enumerate(Y):
        for label_id, label_start, label_end in y:
            labels[i, label_start] = label_id
            output_attn[i, label_start, label_start:label_end] = 1.0 / (label_end.item() - label_start.item())

    batch = {
        'net_input': {
            'src_tokens': X,
            'src_length': lengths,
        },
        'target': labels,
    }
    return batch

def _tokenize(js_tokens, labels, sp, tgt_dict, max_length, split_source_targets_by_tab=False):
    """Given a line from the .txt data files in DeepTyper, format and
    tokenize the code into subwords while retaining a mapping between
    type labels and the subwords.

    Returns:
        Beautified program
        List of subword IDs
        List of (label_id, label_start, label_end) tuples where label_start/end specify a range of subword IDs"""
    assert TYPED_MARKER_START not in js_tokens
    assert TYPED_MARKER_MID not in js_tokens
    assert TYPED_MARKER_END not in js_tokens
    cap_length = max_length != -1
    #     # Code tokens and type labels are delimeted by tab, as in .json files
    #     # Code tokens and type labels are delimited by space after </s>, as in .txt files

    # Split code by spaces to get DeepTyper tokens, excluding <s>, </s>
    js_tokens = js_tokens.split(" ")[1:-1]
    labels = labels.split(" ")[1:-1]
    assert len(js_tokens) == len(labels)

    # Add markers to each labeled token: tokens without no-type label
    js_tokens_with_markers = []
    for tok, label in zip(js_tokens, labels):
        if label != "O":
            tok = f"{TYPED_MARKER_START}{tok}{TYPED_MARKER_END}"
        js_tokens_with_markers.append(tok)

    # Normalize program by beautifying
    js_beautified = jsbeautifier.beautify(" ".join(js_tokens_with_markers))
    js_beautified_norm = normalize_program(js_beautified)
    # js_beautified_norm = js_beautified_norm

    # Subword tokenize, separately tokenizing each marked identifier
    js_buffer = js_beautified_norm
    subword_ids = [sp.PieceToId("<s>")]
    label_segments = []

    last_identifier_end = 0
    start = js_buffer.find(TYPED_MARKER_START)
    labels = list(filter(lambda l: l != "O", labels))
    label_i = 0
    if start < 0:
        # No labeled words in this text, just tokenize the full program
        buffer_ids = sp.EncodeAsIds(js_buffer)
        subword_ids.extend(buffer_ids[: max_length - 2] if cap_length else buffer_ids)
    while start >= 0:
        # Tokenize stuff before the typed identifier
        buffer_ids = sp.EncodeAsIds(js_buffer[last_identifier_end:start])
        if cap_length and len(subword_ids) + len(buffer_ids) + 1 > max_length:  # +1 for </s>
            break
        subword_ids.extend(buffer_ids)

        # Tokenize typed identifier and record label
        start = start + len(TYPED_MARKER_START)
        end = js_buffer.index(TYPED_MARKER_END, start)
        assert end > start, "Zero-length identifier"
        identifier = js_buffer[start:end]
        identifier_ids = sp.EncodeAsIds(identifier)
        if cap_length and len(subword_ids) + len(identifier_ids) + 1 > max_length:  # +1 for </s>
            break
        label_id = tgt_dict.indices.get(labels[label_i], tgt_dict.indices['$any$'])

        label_segments.append((label_id, len(subword_ids), len(subword_ids) + len(identifier_ids)))
        subword_ids.extend(identifier_ids)

        start = js_buffer.find(TYPED_MARKER_START, start + 1)
        last_identifier_end = end + len(TYPED_MARKER_END)
        label_i += 1
    subword_ids.append(sp.PieceToId("</s>"))
    assert subword_ids[0] == sp.PieceToId("<s>")
    assert subword_ids[-1] == sp.PieceToId("</s>")

    return js_beautified, subword_ids, label_segments

class CodeTypeDataset(NccDataset):
    """
    A pair of torch.utils.data.Datasets.

    Args:
        src (torch.utils.data.Dataset): source dataset to wrap
        src_sizes (List[int]): source sentence lengths
        src_dict (~fairseq.data.Dictionary): source vocabulary
        tgt (torch.utils.data.Dataset, optional): target dataset to wrap
        tgt_sizes (List[int], optional): target sentence lengths
        tgt_dict (~fairseq.data.Dictionary, optional): target vocabulary
        left_pad_source (bool, optional): pad source tensors on the left side
            (default: True).
        left_pad_target (bool, optional): pad target tensors on the left side
            (default: False).
        max_source_positions (int, optional): max number of tokens in the
            source sentence (default: 1024).
        max_target_positions (int, optional): max number of tokens in the
            target sentence (default: 1024).
        shuffle (bool, optional): shuffle dataset elements before batching
            (default: True).
        input_feeding (bool, optional): create a shifted version of the targets
            to be passed into the model for teacher forcing (default: True).
        remove_eos_from_source (bool, optional): if set, removes eos from end
            of source if it's present (default: False).
        append_eos_to_target (bool, optional): if set, appends eos to end of
            target if it's absent (default: False).
        align_dataset (torch.utils.data.Dataset, optional): dataset
            containing alignments.
        append_bos (bool, optional): if set, appends bos to the beginning of
            source/target sentence.
    """

    def __init__(
            self, src, src_sizes, src_dict, tgt, tgt_dict, sp,
            left_pad_source=False, left_pad_target=False,
            max_source_positions=1024,
            max_target_positions=1024,
            shuffle=True,
    ):
        self.src = src
        self.src_sizes = np.array(src_sizes)
        self.src_dict = src_dict
        self.tgt = tgt
        self.tgt_dict = tgt_dict
        self.sp = sp
        self.left_pad_source = left_pad_source
        self.left_pad_target = left_pad_target
        self.max_source_positions = max_source_positions
        self.max_target_positions = max_target_positions
        self.shuffle = shuffle

    def __getitem__(self, index):
        # Append EOS to end of tgt sentence if it does not have an EOS
        # and remove EOS from end of src sentence if it exists.
        # This is useful when we use existing datasets for opposite directions
        #   i.e., when we want to use tgt_dataset as src_dataset and vice versa
        src_item = self.src[index]
        tgt_item = self.tgt[index]

        _, subword_ids, label_segments = _tokenize(
            src_item, tgt_item, self.sp, self.tgt_dict, self.max_source_positions)
        if self.max_source_positions != -1:
            assert len(subword_ids) <= self.max_source_positions
        subword_ids = torch.tensor(subword_ids, dtype=torch.long)
        label_segments = torch.tensor(label_segments, dtype=torch.long)

        example = {
            'subword_ids': subword_ids,
            'label_segments': label_segments,
        }

        return example

    def __len__(self):
        return len(self.tgt)

    def collater(self, samples):
        """Merge a list of samples to form a mini-batch.

        Args:
            samples (List[dict]): samples to collate

        Returns:
            dict: a mini-batch with the following keys:

                - `id` (LongTensor): example IDs in the original input order
                - `ntokens` (int): total number of tokens in the batch
                - `net_input` (dict): the input to the Model, containing keys:

                  - `src_tokens` (LongTensor): a padded 2D Tensor of tokens in
                    the source sentence of shape `(bsz, src_len)`. Padding will
                    appear on the left if *left_pad_source* is ``True``.
                  - `src_lengths` (LongTensor): 1D Tensor of the unpadded
                    lengths of each source sentence of shape `(bsz)`
                  - `prev_output_tokens` (LongTensor): a padded 2D Tensor of
                    tokens in the target sentence, shifted right by one
                    position for teacher forcing, of shape `(bsz, tgt_len)`.
                    This key will not be present if *input_feeding* is
                    ``False``.  Padding will appear on the left if
                    *left_pad_target* is ``True``.

                - `target` (LongTensor): a padded 2D Tensor of tokens in the
                  target sentence of shape `(bsz, tgt_len)`. Padding will appear
                  on the left if *left_pad_target* is ``True``.
        """
        return collate(
            samples, pad_idx=self.tgt_dict.pad(), no_type_id=self.tgt_dict.index('O'),
        )

    def num_tokens(self, index):
        """Return the number of tokens in a sample. This value is used to
        enforce ``--max-tokens`` during batching."""
        return self.src_sizes[index]

    def size(self, index):

    def ordered_indices(self):
        """Return an ordered list of indices. Batches will be constructed based
        on this order."""
        if self.shuffle:
            indices = np.random.permutation(len(self))
        else:
            indices = np.arange(len(self))
        return indices#[np.argsort(self.src_sizes[indices], kind='mergesort')] # TODO: debug


        return collate(
            samples, pad_idx=self.tgt_dict.pad(), no_type_id=self.tgt_dict.index('O'),
        )

import numpy as np
from ncc.data.tools import data_utils
from ncc.data.ncc_dataset import NccDataset
import jsbeautifier
from dataset.augmented_javascript.utils.util import normalize_program
import torch

TYPED_MARKER_START = "__LS__"
TYPED_MARKER_MID = "__LM__"
TYPED_MARKER_END = "__LE__"

def collate(samples, pad_idx, no_type_id):
    if len(samples) == 0:
        return {}

    def merge(key):
        return data_utils.collate_tokens(
            [s[key] for s in samples],
            pad_idx,
        )
    B = len(samples)
    X = merge('subword_ids')
    Y = [s['label_segments'] for s in samples]

    L = X.size(1)

    # Create tensor of sequence lengths, [B]
    lengths = torch.LongTensor([s['subword_ids'].numel() for s in samples])
    # Make masks for each label interval
    labels = torch.zeros(B, L, dtype=torch.long)
    labels.fill_(no_type_id)
    output_attn = torch.zeros(B, L, L, dtype=torch.float)
    for i, y in enumerate(Y):
        for label_id, label_start, label_end in y:
            labels[i, label_start] = label_id
            output_attn[i, label_start, label_start:label_end] = 1.0 / (label_end.item() - label_start.item())

    batch = {
        'net_input': {
            'src_tokens': X,
            'src_length': lengths,
        },
        'target': labels,
    }
    return batch

def _tokenize(js_tokens, labels, sp, tgt_dict, max_length, split_source_targets_by_tab=False):
    """Given a line from the .txt data files in DeepTyper, format and
    tokenize the code into subwords while retaining a mapping between
    type labels and the subwords.

    Returns:
        Beautified program
        List of subword IDs
        List of (label_id, label_start, label_end) tuples where label_start/end specify a range of subword IDs"""
    assert TYPED_MARKER_START not in js_tokens
    assert TYPED_MARKER_MID not in js_tokens
    assert TYPED_MARKER_END not in js_tokens
    cap_length = max_length != -1
    #     # Code tokens and type labels are delimeted by tab, as in .json files
    #     # Code tokens and type labels are delimited by space after </s>, as in .txt files

    # Split code by spaces to get DeepTyper tokens, excluding <s>, </s>
    js_tokens = js_tokens.split(" ")[1:-1]
    labels = labels.split(" ")[1:-1]
    assert len(js_tokens) == len(labels)

    # Add markers to each labeled token: tokens without no-type label
    js_tokens_with_markers = []
    for tok, label in zip(js_tokens, labels):
        if label != "O":
            tok = f"{TYPED_MARKER_START}{tok}{TYPED_MARKER_END}"
        js_tokens_with_markers.append(tok)

    # Normalize program by beautifying
    js_beautified = jsbeautifier.beautify(" ".join(js_tokens_with_markers))
    js_beautified_norm = normalize_program(js_beautified)
    # js_beautified_norm = js_beautified_norm

    # Subword tokenize, separately tokenizing each marked identifier
    js_buffer = js_beautified_norm
    subword_ids = [sp.PieceToId("<s>")]
    label_segments = []

    last_identifier_end = 0
    start = js_buffer.find(TYPED_MARKER_START)
    labels = list(filter(lambda l: l != "O", labels))
    label_i = 0
    if start < 0:
        # No labeled words in this text, just tokenize the full program
        buffer_ids = sp.EncodeAsIds(js_buffer)
        subword_ids.extend(buffer_ids[: max_length - 2] if cap_length else buffer_ids)
    while start >= 0:
        # Tokenize stuff before the typed identifier
        buffer_ids = sp.EncodeAsIds(js_buffer[last_identifier_end:start])
        if cap_length and len(subword_ids) + len(buffer_ids) + 1 > max_length:  # +1 for </s>
            break
        subword_ids.extend(buffer_ids)

        # Tokenize typed identifier and record label
        start = start + len(TYPED_MARKER_START)
        end = js_buffer.index(TYPED_MARKER_END, start)
        assert end > start, "Zero-length identifier"
        identifier = js_buffer[start:end]
        identifier_ids = sp.EncodeAsIds(identifier)
        if cap_length and len(subword_ids) + len(identifier_ids) + 1 > max_length:  # +1 for </s>
            break
        label_id = tgt_dict.indices.get(labels[label_i], tgt_dict.indices['$any$'])

        label_segments.append((label_id, len(subword_ids), len(subword_ids) + len(identifier_ids)))
        subword_ids.extend(identifier_ids)

        start = js_buffer.find(TYPED_MARKER_START, start + 1)
        last_identifier_end = end + len(TYPED_MARKER_END)
        label_i += 1
    subword_ids.append(sp.PieceToId("</s>"))
    assert subword_ids[0] == sp.PieceToId("<s>")
    assert subword_ids[-1] == sp.PieceToId("</s>")

    return js_beautified, subword_ids, label_segments

class CodeTypeDataset(NccDataset):
    """
    A pair of torch.utils.data.Datasets.

    Args:
        src (torch.utils.data.Dataset): source dataset to wrap
        src_sizes (List[int]): source sentence lengths
        src_dict (~fairseq.data.Dictionary): source vocabulary
        tgt (torch.utils.data.Dataset, optional): target dataset to wrap
        tgt_sizes (List[int], optional): target sentence lengths
        tgt_dict (~fairseq.data.Dictionary, optional): target vocabulary
        left_pad_source (bool, optional): pad source tensors on the left side
            (default: True).
        left_pad_target (bool, optional): pad target tensors on the left side
            (default: False).
        max_source_positions (int, optional): max number of tokens in the
            source sentence (default: 1024).
        max_target_positions (int, optional): max number of tokens in the
            target sentence (default: 1024).
        shuffle (bool, optional): shuffle dataset elements before batching
            (default: True).
        input_feeding (bool, optional): create a shifted version of the targets
            to be passed into the model for teacher forcing (default: True).
        remove_eos_from_source (bool, optional): if set, removes eos from end
            of source if it's present (default: False).
        append_eos_to_target (bool, optional): if set, appends eos to end of
            target if it's absent (default: False).
        align_dataset (torch.utils.data.Dataset, optional): dataset
            containing alignments.
        append_bos (bool, optional): if set, appends bos to the beginning of
            source/target sentence.
    """

    def __init__(
            self, src, src_sizes, src_dict, tgt, tgt_dict, sp,
            left_pad_source=False, left_pad_target=False,
            max_source_positions=1024,
            max_target_positions=1024,
            shuffle=True,
    ):
        self.src = src
        self.src_sizes = np.array(src_sizes)
        self.src_dict = src_dict
        self.tgt = tgt
        self.tgt_dict = tgt_dict
        self.sp = sp
        self.left_pad_source = left_pad_source
        self.left_pad_target = left_pad_target
        self.max_source_positions = max_source_positions
        self.max_target_positions = max_target_positions
        self.shuffle = shuffle

    def __getitem__(self, index):
        # Append EOS to end of tgt sentence if it does not have an EOS
        # and remove EOS from end of src sentence if it exists.
        # This is useful when we use existing datasets for opposite directions
        #   i.e., when we want to use tgt_dataset as src_dataset and vice versa
        src_item = self.src[index]
        tgt_item = self.tgt[index]

        _, subword_ids, label_segments = _tokenize(
            src_item, tgt_item, self.sp, self.tgt_dict, self.max_source_positions)
        if self.max_source_positions != -1:
            assert len(subword_ids) <= self.max_source_positions
        subword_ids = torch.tensor(subword_ids, dtype=torch.long)
        label_segments = torch.tensor(label_segments, dtype=torch.long)

        example = {
            'subword_ids': subword_ids,
            'label_segments': label_segments,
        }

        return example

    def __len__(self):
        return len(self.tgt)

    def collater(self, samples):
        """Merge a list of samples to form a mini-batch.

        Args:
            samples (List[dict]): samples to collate

        Returns:
            dict: a mini-batch with the following keys:

                - `id` (LongTensor): example IDs in the original input order
                - `ntokens` (int): total number of tokens in the batch
                - `net_input` (dict): the input to the Model, containing keys:

                  - `src_tokens` (LongTensor): a padded 2D Tensor of tokens in
                    the source sentence of shape `(bsz, src_len)`. Padding will
                    appear on the left if *left_pad_source* is ``True``.
                  - `src_lengths` (LongTensor): 1D Tensor of the unpadded
                    lengths of each source sentence of shape `(bsz)`
                  - `prev_output_tokens` (LongTensor): a padded 2D Tensor of
                    tokens in the target sentence, shifted right by one
                    position for teacher forcing, of shape `(bsz, tgt_len)`.
                    This key will not be present if *input_feeding* is
                    ``False``.  Padding will appear on the left if
                    *left_pad_target* is ``True``.

                - `target` (LongTensor): a padded 2D Tensor of tokens in the
                  target sentence of shape `(bsz, tgt_len)`. Padding will appear
                  on the left if *left_pad_target* is ``True``.
        """

    def num_tokens(self, index):
        """Return the number of tokens in a sample. This value is used to
        enforce ``--max-tokens`` during batching."""
        return self.src_sizes[index]

    def size(self, index):
        """Return an example's size as a float or tuple. This value is used when
        filtering a dataset with ``--max-positions``."""
        return self.src_sizes[index]

    def ordered_indices(self):
        """Return an ordered list of indices. Batches will be constructed based
        on this order."""
        if self.shuffle:
            indices = np.random.permutation(len(self))
        else:
            indices = np.arange(len(self))
        return indices#[np.argsort(self.src_sizes[indices], kind='mergesort')] # TODO: debug