__docformat__ = 'restructuredtext'

import sys

from zope import component
from zope.interface import implements
from zope.component.interfaces import ComponentLookupError
from zope.traversing.interfaces import IPathAdapter, ITraversable
from zope.traversing.interfaces import TraversalError
from zope.traversing.adapters import traversePathElement
from zope.security.untrustedpython import rcompile
from zope.security.proxy import ProxyFactory, removeSecurityProxy
from zope.security.untrustedpython.builtins import SafeBuiltins
from zope.i18n import translate

from zope.tales.expressions import PathExpr, StringExpr, NotExpr, DeferExpr
from zope.tales.expressions import SimpleModuleImporter
from zope.tales.pythonexpr import PythonExpr
from zope.tales.tales import ExpressionEngine, Context

from i18n import ZopeMessageFactory as _

class InlineCodeError(Exception):
    pass

class ZopeTraverser(object):

    def __init__(self, proxify=None):
        if proxify is None:
            self.proxify = lambda x: x
        else:
            self.proxify = proxify

    def __call__(self, object, path_items, econtext):
        """Traverses a sequence of names, first trying attributes then items.
        """
        request = getattr(econtext, 'request', None)
        path_items = list(path_items)
        path_items.reverse()

        while path_items:
            name = path_items.pop()

            # special-case dicts for performance reasons
            if getattr(object, '__class__', None) == dict:
                object = object[name]
            else:
                object = traversePathElement(object, name, path_items,
                                             request=request)
            object = self.proxify(object)
        return object

zopeTraverser = ZopeTraverser(ProxyFactory)

class ZopePathExpr(PathExpr):

    def __init__(self, name, expr, engine):
        super(ZopePathExpr, self).__init__(name, expr, engine, zopeTraverser)

trustedZopeTraverser = ZopeTraverser()

class TrustedZopePathExpr(PathExpr):

    def __init__(self, name, expr, engine):
        super(TrustedZopePathExpr, self).__init__(name, expr, engine,
                                                  trustedZopeTraverser)

# Create a version of the restricted built-ins that uses a safe
# version of getattr() that wraps values in security proxies where
# appropriate:

class ZopePythonExpr(PythonExpr):

    def __call__(self, econtext):
        __traceback_info__ = self.text
        vars = self._bind_used_names(econtext, SafeBuiltins)
        return eval(self._code, vars)

    def _compile(self, text, filename):
        return rcompile.compile(text, filename, 'eval')

class ZopeContextBase(Context):
    """Base class for both trusted and untrusted evaluation contexts."""

    def translate(self, msgid, domain=None, mapping=None, default=None):
        return translate(msgid, domain, mapping,
                         context=self.request, default=default)

    evaluateInlineCode = False

    def evaluateCode(self, lang, code):
        if not self.evaluateInlineCode:
            raise InlineCodeError(
                  _('Inline Code Evaluation is deactivated, which means that '
                    'you cannot have inline code snippets in your Page '
                    'Template. Activate Inline Code Evaluation and try again.'))

        # TODO This is only needed when self.evaluateInlineCode is true,
        # so should only be needed for zope.app.pythonpage.
        from zope.app.interpreter.interfaces import IInterpreter
        interpreter = component.queryUtility(IInterpreter, lang)
        if interpreter is None:
            error = _('No interpreter named "${lang_name}" was found.',
                      mapping={'lang_name': lang})
            raise InlineCodeError(error)

        globals = self.vars.copy()
        result = interpreter.evaluateRawCode(code, globals)
        # Add possibly new global variables.
        old_names = self.vars.keys()
        for name, value in globals.items():
            if name not in old_names:
                self.setGlobal(name, value)
        return result

class ZopeContext(ZopeContextBase):
    """Evaluation context for untrusted programs."""

    def evaluateMacro(self, expr):
        """evaluateMacro gets security-proxied macro programs when this
        is run with the zopeTraverser, and in other untrusted
        situations. This will cause evaluation to fail in
        zope.tal.talinterpreter, which knows nothing of security proxies.
        Therefore, this method removes any proxy from the evaluated
        expression.

        >>> output = [('version', 'xxx'), ('mode', 'html'), ('other', 'things')]
        >>> def expression(context):
        ...     return ProxyFactory(output)
        ...
        >>> zc = ZopeContext(ExpressionEngine, {})
        >>> out = zc.evaluateMacro(expression)
        >>> type(out)
        <type 'list'>

        The method does some trivial checking to make sure we are getting
        back a macro like we expect: it must be a sequence of sequences, in
        which the first sequence must start with 'version', and the second
        must start with 'mode'.

        >>> del output[0]
        >>> zc.evaluateMacro(expression) # doctest: +ELLIPSIS
        Traceback (most recent call last):
        ...
        ValueError: ('unexpected result from macro evaluation.', ...)

        >>> del output[:]
        >>> zc.evaluateMacro(expression) # doctest: +ELLIPSIS
        Traceback (most recent call last):
        ...
        ValueError: ('unexpected result from macro evaluation.', ...)

        >>> output = None
        >>> zc.evaluateMacro(expression) # doctest: +ELLIPSIS
        Traceback (most recent call last):
        ...
        ValueError: ('unexpected result from macro evaluation.', ...)
        """
        macro = removeSecurityProxy(Context.evaluateMacro(self, expr))
        # we'll do some basic checks that it is the sort of thing we expect
        problem = False
        try:
            problem = macro[0][0] != 'version' or macro[1][0] != 'mode'
        except (TypeError, IndexError):
            problem = True
        if problem:
            raise ValueError('unexpected result from macro evaluation.', macro)
        return macro

    def setContext(self, name, value):
        # Hook to allow subclasses to do things like adding security proxies
        Context.setContext(self, name, ProxyFactory(value))

class TrustedZopeContext(ZopeContextBase):
    """Evaluation context for trusted programs."""

class AdapterNamespaces(object):
    """Simulate tales function namespaces with adapter lookup.

    When we are asked for a namespace, we return an object that
    actually computes an adapter when called:

    To demonstrate this, we need to register an adapter:

      >>> from zope.component.testing import setUp, tearDown
      >>> setUp()
      >>> from zope.component import provideAdapter
      >>> def adapter1(ob):
      ...     return 1
      >>> adapter1.__component_adapts__ = (None,)
      >>> provideAdapter(adapter1, None, IPathAdapter, 'a1')

    Now, with this adapter in place, we can try out the namespaces:

      >>> ob = object()
      >>> namespaces = AdapterNamespaces()
      >>> namespace = namespaces['a1']
      >>> namespace(ob)
      1
      >>> namespace = namespaces['a2']
      >>> namespace(ob)
      Traceback (most recent call last):
      ...
      KeyError: 'a2'

    Cleanup:

      >>> tearDown()
    """

    def __init__(self):
        self.namespaces = {}

    def __getitem__(self, name):
        namespace = self.namespaces.get(name)
        if namespace is None:
            def namespace(object):
                try:
                    return component.getAdapter(object, IPathAdapter, name)
                except ComponentLookupError:
                    raise KeyError(name)

            self.namespaces[name] = namespace
        return namespace

class ZopeBaseEngine(ExpressionEngine):

    _create_context = ZopeContext

    def __init__(self):
        ExpressionEngine.__init__(self)
        self.namespaces = AdapterNamespaces()

    def getContext(self, __namespace=None, **namespace):
        if __namespace:
            if namespace:
                namespace.update(__namespace)
            else:
                namespace = __namespace

        context = self._create_context(self, namespace)

        # Put request into context so path traversal can find it
        if 'request' in namespace:
            context.request = namespace['request']

        # Put context into context so path traversal can find it
        if 'context' in namespace:
            context.context = namespace['context']

        return context

class ZopeEngine(ZopeBaseEngine):
    """Untrusted expression engine.

    This engine does not allow modules to be imported; only modules
    already available may be accessed::

      >>> modname = 'zope.pagetemplate.tests.trusted'
      >>> engine = _Engine()
      >>> context = engine.getContext(engine.getBaseNames())

      >>> modname in sys.modules
      False
      >>> context.evaluate('modules/' + modname)
      Traceback (most recent call last):
        ...
      KeyError: 'zope.pagetemplate.tests.trusted'

    (The use of ``KeyError`` is an unfortunate implementation detail; I
    think this should be a ``TraversalError``.)

    Modules which have already been imported by trusted code are
    available, wrapped in security proxies::

      >>> m = context.evaluate('modules/sys')
      >>> m.__name__
      'sys'
      >>> m._getframe
      Traceback (most recent call last):
        ...
      ForbiddenAttribute: ('_getframe', <module 'sys' (built-in)>)

    The results of Python expressions evaluated by this engine are
    wrapped in security proxies::

      >>> r = context.evaluate('python: {12: object()}.values')
      >>> type(r)
      <type 'zope.security._proxy._Proxy'>
      >>> r = context.evaluate('python: {12: object()}.values()[0].__class__')
      >>> type(r)
      <type 'zope.security._proxy._Proxy'>

    General path expressions provide objects that are wrapped in
    security proxies as well::

      >>> from zope.component.testing import setUp, tearDown
      >>> from zope.security.checker import NamesChecker, defineChecker

      >>> class Container(dict):
      ...     implements(ITraversable)
      ...     def traverse(self, name, further_path):
      ...         return self[name]

      >>> setUp()
      >>> defineChecker(Container, NamesChecker(['traverse']))
      >>> d = engine.getBaseNames()
      >>> foo = Container()
      >>> foo.__name__ = 'foo'
      >>> d['foo'] = ProxyFactory(foo)
      >>> foo['bar'] = bar = Container()
      >>> bar.__name__ = 'bar'
      >>> bar.__parent__ = foo
      >>> bar['baz'] = baz = Container()
      >>> baz.__name__ = 'baz'
      >>> baz.__parent__ = bar
      >>> context = engine.getContext(d)

      >>> o1 = context.evaluate('foo/bar')
      >>> o1.__name__
      'bar'
      >>> type(o1)
      <type 'zope.security._proxy._Proxy'>

      >>> o2 = context.evaluate('foo/bar/baz')
      >>> o2.__name__
      'baz'
      >>> type(o2)
      <type 'zope.security._proxy._Proxy'>
      >>> o3 = o2.__parent__
      >>> type(o3)
      <type 'zope.security._proxy._Proxy'>
      >>> o1 == o3
      True

      >>> o1 is o2
      False

    Note that this engine special-cases dicts during path traversal:
    it traverses only to their items, but not to their attributes
    (e.g. methods on dicts), because of performance reasons:

      >>> d = engine.getBaseNames()
      >>> d['adict'] = {'items': 123}
      >>> d['anotherdict'] = {}
      >>> context = engine.getContext(d)
      >>> context.evaluate('adict/items')
      123
      >>> context.evaluate('anotherdict/keys')
      Traceback (most recent call last):
        ...
      KeyError: 'keys'

      >>> tearDown()

    """

    def getFunctionNamespace(self, namespacename):
        """ Returns the function namespace """
        return ProxyFactory(
            super(ZopeEngine, self).getFunctionNamespace(namespacename))

class TrustedZopeEngine(ZopeBaseEngine):
    """Trusted expression engine.

    This engine allows modules to be imported::

      >>> modname = 'zope.pagetemplate.tests.trusted'
      >>> engine = _TrustedEngine()
      >>> context = engine.getContext(engine.getBaseNames())

      >>> modname in sys.modules
      False
      >>> m = context.evaluate('modules/' + modname)
      >>> m.__name__ == modname
      True
      >>> modname in sys.modules
      True

    Since this is trusted code, we can look at whatever is in the
    module, not just __name__ or what's declared in a security
    assertion::

      >>> m.x
      42

    Clean up after ourselves::

      >>> del sys.modules[modname]

    """

    _create_context = TrustedZopeContext

class TraversableModuleImporter(SimpleModuleImporter):

    implements(ITraversable)

    def traverse(self, name, further_path):
        try:
            return self[name]
        except KeyError:
            raise TraversalError(self, name)

def _Engine(engine=None):
    if engine is None:
        engine = ZopeEngine()
    engine = _create_base_engine(engine, ZopePathExpr)
    engine.registerType('python', ZopePythonExpr)

    # Using a proxy around sys.modules allows page templates to use
    # modules for which security declarations have been made, but
    # disallows execution of any import-time code for modules, which
    # should not be allowed to happen during rendering.
    engine.registerBaseName('modules', ProxyFactory(sys.modules))

    return engine

def _TrustedEngine(engine=None):
    if engine is None:
        engine = TrustedZopeEngine()
    engine = _create_base_engine(engine, TrustedZopePathExpr)
    engine.registerType('python', PythonExpr)
    engine.registerBaseName('modules', TraversableModuleImporter())
    return engine

def _create_base_engine(engine, pathtype):
    for pt in pathtype._default_type_names:
        engine.registerType(pt, pathtype)
    engine.registerType('string', StringExpr)
    engine.registerType('not', NotExpr)
    engine.registerType('defer', DeferExpr)
    return engine

Engine = _Engine()
TrustedEngine = _TrustedEngine()

class AppPT(object):

    def pt_getEngine(self):
        return Engine

class TrustedAppPT(object):

    def pt_getEngine(self):
        return TrustedEngine


'''
https://leetcode.com/problems/minimum-height-trees/description/
'''

from collections import defaultdict, deque

# there can only be at most two roots in MHT, which are the mid points of the longest path
# find leaves, remove leaves from neighbours
# keep doing as long as num of nodes is > 2
# ones remaining are the root
def findMinHeightTrees(n, edges):
    """
    :type n: int
    :type edges: List[List[int]]
    :rtype: List[int]
    """
    adjlist = defaultdict(set)
    for u,v in edges:
        adjlist[u].add(v)
        adjlist[v].add(u)

    nodes = set(range(n)) # {0, 1, 2, 3}
    while len(nodes) > 2:
        leaves = {i for i in nodes if len(adjlist[i]) == 1}
        nodes -= leaves # remove set of leaves from nodes
        for i in leaves:
            for edge in adjlist[i]: # moving up one level
                adjlist[edge].remove(i)
    return list(nodes)

print(findMinHeightTrees(6, [[0, 3], [1, 3], [2, 3], [4, 3], [5, 4]])) # [3, 4]

# runs faster due to one lesser for loop, and no set comprehension
def find_min_height_faster(n, edges):
    adjlist = defaultdict(set)
    for u,v in edges:
        adjlist[u].add(v)
        adjlist[v].add(u)
    nodes = set(range(n)) # {0, 1, 2, 3}
    leaves = {i for i in nodes if len(adjlist[i]) == 1}

    while len(nodes) > 2:
        newleaves = set()
        nodes -= leaves
        for i in leaves:
            j = adjlist[i].pop()
            adjlist[j].remove(i)
            print(adjlist)
            if len(adjlist[j]) == 1:
                newleaves.add(j)
        leaves = newleaves
    return list(nodes)

# bfs, O(n^2), TLE
# do bfs on each node to find minheight of each node
def findMinHeightTrees(n, edges):
    """
    :type n: int
    :type edges: List[List[int]]
    :rtype: List[int]
    """
    adjlist = defaultdict(list)
    for u,v in edges:
        adjlist[u].append(v)
        adjlist[v].append(u)

    levels, ans, minheight = {}, [], float('inf')
    for i in range(n):
        visited, queue = set(), deque([(i, 0)])
        level, height = 0, float('-inf')

        while queue:
            curr, level = queue.popleft()
            visited.add(curr)
            for edge in adjlist[curr]:
                if edge not in visited:
                    queue.append((edge, level+1))
        levels[i] = level
        minheight = min(minheight, level)

    for key, val in levels.items():
        if val == minheight:
            ans.append(key)
    return ans

print(findMinHeightTrees(4, [[1, 0], [1, 2], [1, 3]])) # [1]


from __future__ import annotations

import os
import typing as t
import logging
import pathlib
from typing import TYPE_CHECKING

import yaml
import pytest

import bentoml
from bentoml._internal.utils import bentoml_cattr
from bentoml._internal.models import ModelStore
from bentoml._internal.models import ModelContext
from bentoml._internal.bento.build_config import BentoBuildConfig

if TYPE_CHECKING:
    from _pytest.python import Metafunc

TEST_MODEL_CONTEXT = ModelContext(
    framework_name="testing", framework_versions={"testing": "v1"}
)

def pytest_generate_tests(metafunc: Metafunc) -> None:
    from bentoml._internal.utils import analytics

    analytics.usage_stats.do_not_track.cache_clear()
    analytics.usage_stats._usage_event_debugging.cache_clear()  # type: ignore (private warning)

    # used for local testing, on CI we already set DO_NOT_TRACK
    os.environ["__BENTOML_DEBUG_USAGE"] = "False"
    os.environ["BENTOML_DO_NOT_TRACK"] = "True"

@pytest.fixture(scope="function")
def noop_service(dummy_model_store: ModelStore) -> bentoml.Service:
    import cloudpickle

    from bentoml.io import Text

    class NoopModel:
        def predict(self, data: t.Any) -> t.Any:
            return data

    with bentoml.models.create(
        "noop_model",
        context=TEST_MODEL_CONTEXT,
        module=__name__,
        signatures={"predict": {"batchable": True}},
        _model_store=dummy_model_store,
    ) as model:
        with open(model.path_of("test.pkl"), "wb") as f:
            cloudpickle.dump(NoopModel(), f)

    ref = bentoml.models.get("noop_model", _model_store=dummy_model_store)

    class NoopRunnable(bentoml.Runnable):
        SUPPORTED_RESOURCES = ("cpu",)
        SUPPORTS_CPU_MULTI_THREADING = True

        def __init__(self):
            self._model: NoopModel = bentoml.picklable_model.load_model(ref)

        @bentoml.Runnable.method(batchable=True)
        def predict(self, data: t.Any) -> t.Any:
            return self._model.predict(data)

    svc = bentoml.Service(
        name="noop_service",
        runners=[bentoml.Runner(NoopRunnable, models=[ref])],
    )

    @svc.api(input=Text(), output=Text())
    def noop_sync(data: str) -> str:  # type: ignore
        return data

    return svc

@pytest.fixture(scope="function", autouse=True, name="propagate_logs")
def fixture_propagate_logs() -> t.Generator[None, None, None]:
    logger = logging.getLogger("bentoml")
    # bentoml sets propagate to False by default, so we need to set it to True
    # for pytest caplog to recognize logs
    logger.propagate = True
    yield
    # restore propagate to False after tests
    logger.propagate = False

@pytest.fixture(scope="function")
def reload_directory(
    request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory
) -> t.Generator[pathlib.Path, None, None]:
    """
    This fixture will create an example bentoml working file directory
    and yield the results directory
    ./
    ├── models/  # mock default bentoml home models directory
    ├── [fdir, fdir_one, fdir_two]/
    │   ├── README.md
        ├── subdir/
        │   ├── README.md
    │   │   └── app.py
    │   ├── somerust.rs
    │   └── app.py
    ├── README.md
    ├── .bentoignore
    ├── bentofile.yaml
    ├── fname.ipynb
    ├── requirements.txt
    ├── service.py
    └── train.py
    """
    root = tmp_path_factory.mktemp("reload_directory")
    # create a models directory
    root.joinpath("models").mkdir()

    # enable this fixture to use with unittest.TestCase
    if request.cls is not None:
        request.cls.reload_directory = root

    root_file = [
        "README.md",
        "requirements.txt",
        "service.py",
        "train.py",
        "fname.ipynb",
    ]
    for f in root_file:
        p = root.joinpath(f)
        p.touch()

    build_config = BentoBuildConfig(
        service="service.py:svc",
        description="A mock service",
        exclude=["*.rs"],
    ).with_defaults()
    bentofile = root / "bentofile.yaml"
    bentofile.touch()
    with bentofile.open("w", encoding="utf-8") as f:
        yaml.safe_dump(bentoml_cattr.unstructure(build_config), f)

    custom_library = ["fdir", "fdir_one", "fdir_two"]
    for app in custom_library:
        ap = root.joinpath(app)
        ap.mkdir()
        dir_files: list[tuple[str, list[t.Any]]] = [
            ("README.md", []),
            ("subdir", ["README.md", "app.py"]),
            ("lib.rs", []),
            ("app.py", []),
        ]
        for name, maybe_files in dir_files:
            if maybe_files:
                dpath = ap.joinpath(name)
                dpath.mkdir()
                for f in maybe_files:
                    p = dpath.joinpath(f)
                    p.touch()
            else:
                p = ap.joinpath(name)
                p.touch()

    yield root

@pytest.fixture(scope="function", name="change_test_dir")
def fixture_change_test_dir(
    request: pytest.FixtureRequest,
) -> t.Generator[None, None, None]:
    os.chdir(request.fspath.dirname)  # type: ignore (bad pytest stubs)
    yield
    os.chdir(request.config.invocation_dir)  # type: ignore (bad pytest stubs)

@pytest.fixture(scope="session", name="dummy_model_store")
def fixture_dummy_model_store(tmpdir_factory: "pytest.TempPathFactory") -> ModelStore:
    store = ModelStore(tmpdir_factory.mktemp("models"))
    with bentoml.models.create(
        "testmodel",
        module=__name__,
        signatures={},
        context=TEST_MODEL_CONTEXT,
        _model_store=store,
    ):
        pass
    with bentoml.models.create(
        "testmodel",
        module=__name__,
        signatures={},
        context=TEST_MODEL_CONTEXT,
        _model_store=store,
    ):
        pass
    with bentoml.models.create(
        "anothermodel",
        module=__name__,
        signatures={},
        context=TEST_MODEL_CONTEXT,
        _model_store=store,
    ):
        pass

    return store


import json
import argparse
import shutil

import nbt
import os

def main():
    parser = argparse.ArgumentParser(description='修改玩家座標')
    parser.add_argument("username", help="Player name", type=str)
    parser.add_argument("dim", help="Target dim", type=int)
    parser.add_argument("x", help="Target x", type=float)
    parser.add_argument("y", help="Target y", type=float)
    parser.add_argument("z", help="Target z", type=float)
    args = parser.parse_args()

    with open('usernamecache.json') as f:
        uuid_to_username = json.load(f)

    target_uuid = None
    for uuid, username in uuid_to_username.items():
        if username == args.username:
            target_uuid = uuid

    print('Player uuid: ' + target_uuid)

    filename = os.path.join('world', 'playerdata', target_uuid + '.dat')
    shutil.copy(filename, filename + '.bak')

    player_dat = nbt.nbt.NBTFile(filename, 'rb')
    player_dat['Dimension'].value = args.dim
    player_dat['Pos'][0].value = args.x
    player_dat['Pos'][1].value = args.y
    player_dat['Pos'][2].value = args.z

    player_dat.write_file(filename)
    print('Done')

if __name__ == '__main__':
    main()


# -*- coding: utf-8 -*-
"""
Created on Wed Apr 24 11:26:01 2019

@author: soura
"""
import hashtag
import query
import username
import json
import decider
import pandas as pd

class Transfer:
    def __init__(self):
        self.file_mapping = {"query":"query_results.json", "username":"username_results.json","hashtag":"hashtag_results.json"}
        self.predictor=decider.Decider(mode="test")
    def set_mode(self, mode):
        self.mode = mode
    def fetch_predict(self, data):
        filename = self.file_mapping[self.mode]
        if self.mode=="query":
            query.fetch_queries(data)
        elif self.mode=="username":
            username.get_tweets(data)
        else:
            hashtag.get_tweets(data)
        data=json.load(open(filename,'r'))
        tweets=[]
        for obj in data:
            tweets.append(obj['tweet'])
        predictions=[self.predictor.predict(tweets[i])[0] for i in range(len(tweets))]
        return pd.DataFrame({"Tweets":tweets,"Prediction":predictions})


import pandas as pd
import rpy2.robjects.packages as rpackages
from rpy2.robjects.vectors import StrVector
import xlsxwriter

def add_feature_metadata(df_target, df_feature_metadata, df_target_fn_col=None):
    """Adds feature metadata columns to df_target.

    Args:
        df_target - a DataFrame; must contain feature names somewhere
        df_feature_metadata - a DataFrame containing feature metadata,
            as loaded by fst.py's self.feature_metadata
        df_target_fn_col - the name of the column in df_target that contains
            the feature names, or None if it's the index of df_target that
            contains the feature names

    Returns:
        a copy of df_target with the feature metadata added
    """
    original_index = df_target.index
    if df_target_fn_col is not None:
        df_target.index = df_target[df_target_fn_col]
    df_joined = df_target.join(df_feature_metadata)
    # add groups for shadow features
    s_group = df_joined['group']
    feature_names = s_group.index.values
    s_non_shadows = pd.Series(
        [not fname.startswith('shadow_') for fname in feature_names],
        index=feature_names)
    s_group = s_group.where(s_non_shadows, 'Shadow')
    df_joined['group'] = s_group
    # handle any remaining NaN groups
    df_joined['group'].fillna('NA')
    # restore index
    df_target.index = original_index

    return df_joined

def create_refuse_xlsx(df_refuse_table, df_feature_metadata, outpath):
    """Creates an Excel spreadsheet containing the REFUSE Table data,
    annotated with the feature metadata, saved to the given outpath."""

    # merge data
    df_complete = add_feature_metadata(df_refuse_table, df_feature_metadata)
    df_complete['Feature Name'] = df_complete.index.values # easier to iterate

    # sort by VIM_mean
    df_complete.sort(columns='VIM_mean', ascending=False, inplace=True)

    # create spreadsheet
    workbook = xlsxwriter.Workbook(outpath)
    worksheet = workbook.add_worksheet()

    # start from first cell (zero-indexed)
    row = 0
    col = 0

    # styles
    title_style = workbook.add_format({
        'font_size': 12,
        'font_color': 'white',
        'bg_color': '#808080'
        })
    relevant_style = workbook.add_format({
        'font_size': 12,
        'bg_color': '#EBF1DE'
        })
    default_style = workbook.add_format({
        'font_size': 12})

    # these columns should always come first
    column_order = ['Feature Name', 'VIM_mean', 'p_val', 'VIM_var', 'group']
    # there may be user-defined columns to follow
    additional_columns = [x for x in df_complete.columns.values
        if x not in column_order]
    # final order
    column_order += additional_columns

    # some columns get user-friendly names
    column_titles = {
        'VIM_mean': 'Rank (higher=better)',
        'p_val': 'Likelihood of Relevance',
        'VIM_var': 'Variance of Rank',
        'group': 'Group Name'}

    for column in column_order:
        row = 0
        title = column
        if column in column_titles:
            title = column_titles[column]
        worksheet.write(row, col, title, title_style)
        for cell in df_complete[column]:
            row += 1
            style = default_style
            if column == 'p_val' and cell > 0.5:
                style = relevant_style
            worksheet.write(row, col, cell, style)
        col += 1

    # set column widths
    worksheet.set_column(0, len(column_order)-1, 30)

    workbook.close()

def check_r_packages():
    """Ensures required R packages are installed."""
    # import R's utility package
    utils = rpackages.importr('utils')
    # R package names
    packnames = ('classInt',
        'minerva', # should already be installed in dockerfile
        'randomForest', #should already be installed in dockerfile
        #'ggplot2', #used by FST commands that generate PDFs, which aren't currently
        #           #used, and this package doesn't install properly due to versioning problems
        'MASS',
        'gplots')

    # R vector of strings
    for x in packnames:
        print('$$$ ensuring installation of R package: ' + str(x))
        if rpackages.isinstalled(x):
            print('$$$  already installed')
        else:
            print('$$$  INSTALLING: ' + x)
            utils.install_packages(StrVector([x]), verbose=True, quiet=False, dependencies=True, repos="https://cran.rstudio.com", method="curl")
            print('$$$  reported as now installed? ' + str(rpackages.isinstalled(x)) + ' for ' + x)

    #R likes to fail silently, so double check.
    failed_installs = list()
    for x in packnames:
        if not rpackages.isinstalled(x):
            failed_installs.append(x)
    if not len(failed_installs) == 0:
        raise Exception('rpackages claims to have installed these packages, but did not: ' + str(failed_installs))


"""Management of per-player info."""
from functools import partial

from PyQt5.QtCore import Qt
from PyQt5.QtGui import QIntValidator, QPalette, QColor
from PyQt5.QtWidgets import (
    QFormLayout,
    QGroupBox,
    QGridLayout,
    QHBoxLayout,
    QLabel,
    QLineEdit,
    QPushButton,
    QWidget,
)

from table.scorer import Phase, dress_value

class Player(QGroupBox):
    """A widget for managing player info."""

    def __init__(self, name, dress_cb, drop_cb):
        """Initialise with name and a callback for board dressing."""
        super().__init__(name)
        self.name = name
        self.setAlignment(Qt.AlignCenter)

        layout = QFormLayout()

        self.count = QLabel("")
        layout.addRow("Count:", self.count)

        self.cards = QLineEdit()
        self.cards.setValidator(QIntValidator())
        layout.addRow("Cards:", self.cards)

        self.dress = QPushButton("Dress")
        self.dress.clicked.connect(partial(dress_cb, self))
        self.drop = QPushButton("Go Out")
        self.drop.clicked.connect(partial(drop_cb, self))
        self.drop.hide()
        buttons = QHBoxLayout()
        buttons.addWidget(self.dress)
        buttons.addWidget(self.drop)
        layout.setLayout(2, QFormLayout.SpanningRole, buttons)

        self.setLayout(layout)

    def set_color(self, color):
        """Set the color palette."""
        for button in (self.dress, self.drop):
            col = QColor(color)
            if not button.isEnabled():
                col.setAlpha(150)
            pal = button.palette()
            pal.setColor(self.backgroundRole(), col)
            pal.setColor(QPalette.ButtonText, col.darker())
            button.setPalette(pal)

    def set_bold_italic(self, bold_italic):
        """Set whether the font should be both bold and italic or not."""
        font = self.font()
        font.setBold(bold_italic)
        font.setItalic(bold_italic)
        self.setFont(font)

    def refresh(self, phase, is_in_game, is_dresser, balance):
        """Refresh the widget based on the game state."""

        # Update player scores
        self.count.setText(str(balance))

        # Clear card counts
        self.cards.setText("")

        # Deduce other flags
        must_dress = is_dresser and (phase == Phase.DRESSING)
        cannot_dress = balance < dress_value()
        out_of_counters = (must_dress and cannot_dress) or (balance < 0)

        # Update widget availability based on the game state
        self.setEnabled(is_in_game)
        self.cards.setEnabled(phase == Phase.SCORING)
        self.dress.setEnabled(must_dress)
        self.drop.setEnabled(out_of_counters)
        if out_of_counters:
            self.drop.show()
        else:
            self.drop.hide()

        # Update colour based on game state
        if not is_in_game:
            self.set_color(Qt.lightGray)
        elif out_of_counters:
            self.set_color(Qt.red)
        elif cannot_dress:
            self.set_color(Qt.yellow)
        elif must_dress:
            self.set_color(Qt.green)
        else:
            self.set_color(Qt.darkGray)

        # Update font based on whether the player is the dresser
        self.set_bold_italic(is_dresser)

    @property
    def cards_left(self):
        """Get the number of cards remaining."""
        try:
            return int(self.cards.text())
        except ValueError:
            return 0

class PlayerPanel(QWidget):
    """A panel of widgets for managing player info."""

    N_ROWS = 4

    def __init__(self, players, dress_cb, drop_cb):
        """Initialise from a list of names and board dressing callback."""
        super().__init__()
        layout = QGridLayout()

        self.players = {}
        for i, name in enumerate(players):
            player = Player(name, dress_cb, drop_cb)
            self.players[name] = player
            row = 2 * (i % self.N_ROWS)
            col = (i // self.N_ROWS)
            layout.addWidget(player, row, col)
            layout.setRowStretch(row + 1, 1)

        self.setLayout(layout)

    def __getitem__(self, key):
        """Return the requested player widget."""
        return self.players[key]

    def __iter__(self):
        """Return iterator through player widgets."""
        return iter(self.players.values())

    def refresh(self, phase, players, dresser, balance):
        """Refresh the player score displays based on the game state."""

        for name, player in self.players.items():
            player.refresh(phase,
                           name in players,
                           name == dresser,
                           balance[name])


import glob
import cv2
import pandas as pd
import openpyxl
from openpyxl.utils.dataframe import dataframe_to_rows
from skimage.metrics import structural_similarity as ssim
import os
import sys

def SearchExtFiles(fpath, ext, lst):
    try:
        flist = os.listdir(fpath)
    except:
        return
    for f in flist:
        next = os.path.join(fpath, f)
        if(os.path.splitext(next)[-1] == ext):
            lst.append(next)
        if os.path.isdir(next):
            SearchExtFiles(next,ext,lst)

# 열 너비 자동 맞춤
def AutoFitColumnSize(worksheet, columns=None, margin=2):
    for i, column_cells in enumerate(worksheet.columns):
        is_ok = False
        if columns == None:
            is_ok = True
        elif isinstance(columns, list) and i in columns:
            is_ok = True

        if is_ok:
            length = max(len(str(cell.value)) for cell in column_cells)
            worksheet.column_dimensions[column_cells[0].column_letter].width = length + margin

def image_compare():
    # 통합 문서 객체 생성
    wb = openpyxl.Workbook()
    # 시트 선택
    ws = wb.active
    ws.append(['index', 'Similarity', 'different image'])
    refLst = []
    SearchExtFiles(sys.argv[1],'.bmp',refLst)   # main 있는 자리에 이미지가 있으면 경로 지정할 때 ./ 으로 놓기  (경로는 edit_Configuration 에서 parameter 바꾸기)
    compLst = []
    SearchExtFiles(sys.argv[2],'.bmp', compLst)

    # 파일 전체 읽기 위해 glob 함수 사용
    ref_img_list = refLst
    compare_img_list = compLst

    for i in range(len(ref_img_list)):

        imageA = cv2.imread(ref_img_list[i])
        imageB = cv2.imread(compare_img_list[i])
        imageC = imageA.copy()
        if (ref_img_list[i] == compare_img_list[i]):
            tempDiff = cv2.subtract(imageA, imageB)

        grayA = cv2.cvtColor(imageA, cv2.COLOR_BGR2GRAY)
        grayB = cv2.cvtColor(imageB, cv2.COLOR_BGR2GRAY)

        (score, diff) = ssim(grayA, grayB, full=True)
        diff = (diff * 255).astype("uint8")
        print(f"Similarity: {score:.5f}")
        if (score == 1):
            print("같음")
        # 다를 땐 엑셀로 이미지 파일 이름 뽑기
        else:
            diff_image = ref_img_list[i]
            print(diff_image)
            df = pd.DataFrame([[score, diff_image]], index=[ws.max_row/2+0.5], columns=['Similarity', 'different image'])
            # df = pd.DataFrame([[score, diff_image]])
            # 시트에 데이터프레임 삽입
            for r in dataframe_to_rows(df, index=True, header=False):  # True 면 변경 가능 False면 변경 불가능
                ws.append(r)
            print("다름")

        # assert 함수는 디버깅 모드에서 개발자가 오류가 생기면 치명적일 것이라는 곳에 심어 놓는 에러 검출용 코드
        assert score, "다른 점 찾을 수 없음"

    df = df.reset_index(inplace=True)

    AutoFitColumnSize(ws)
    # df.index=df.index+1
    # 엑셀 파일에 저장
    wb.save("./test.xlsx")

    # 아무 키나 눌러서 종료하기
    cv2.waitKey(0)

# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    image_compare()


# -*- coding: utf-8 -*-

#  python-holidays
#  ---------------
#  A fast, efficient Python library for generating country, province and state
#  specific sets of holidays on the fly. It aims to make determining whether a
#  specific date is a holiday as fast and flexible as possible.
#
#  Author:  ryanss <ryanssdev@icloud.com> (c) 2014-2017
#           dr-prodigy <maurizio.montel@gmail.com> (c) 2017-2020
#  Website: https://github.com/dr-prodigy/python-holidays
#  License: MIT (see LICENSE file)

from datetime import date

from dateutil.easter import easter
from dateutil.relativedelta import relativedelta as rd, TH, FR, SA, SU

from holidays.constants import JAN, FEB, MAR, APR, MAY, SEP, OCT, \
    DEC
from holidays.holiday_base import HolidayBase

class Honduras(HolidayBase):
    # https://www.timeanddate.com/holidays/honduras/

    def __init__(self, **kwargs):
        self.country = "HND"
        HolidayBase.__init__(self, **kwargs)

    def _populate(self, year):
        # New Year's Day
        if self.observed and date(year, JAN, 1):
            self[date(year, JAN, 1)] = "Año Nuevo [New Year's Day]"

        # The Three Wise Men Day
        if self.observed and date(year, JAN, 6):
            name = "Día de los Reyes Magos [The Three Wise Men Day] (Observed)"
            self[date(year, JAN, 6)] = name

        # The Three Wise Men Day
        if self.observed and date(year, FEB, 3):
            name = "Día de la virgen de Suyapa [Our Lady of Suyapa] (Observed)"
            self[date(year, FEB, 3)] = name

        # The Father's Day
        if self.observed and date(year, MAR, 19):
            name = "Día del Padre [Father's Day] (Observed)"
            self[date(year, MAR, 19)] = name

        # Maundy Thursday
        self[easter(year) + rd(weekday=TH(-1))
             ] = "Jueves Santo [Maundy Thursday]"

        # Good Friday
        self[easter(year) + rd(weekday=FR(-1))
             ] = "Viernes Santo [Good Friday]"

        # Holy Saturday
        self[easter(year) + rd(weekday=SA(-1))
             ] = "Sábado de Gloria [Holy Saturday]"

        # Easter Sunday
        self[easter(year) + rd(weekday=SU(-1))
             ] = "Domingo de Resurrección [Easter Sunday]"

        # America Day
        if self.observed and date(year, APR, 14):
            self[date(year, APR, 14)] = "Día de las Américas [America Day]"

        # Labor Day
        if self.observed and date(year, MAY, 1):
            self[date(year, MAY, 1)] = "Día del Trabajo [Labour Day]"

        # Mother's Day
        may_first = date(int(year), 5, 1)
        weekday_seq = may_first.weekday()
        mom_day = (14 - weekday_seq)
        if self.observed and date(year, MAY, mom_day):
            str_day = "Día de la madre [Mother's Day] (Observed)"
            self[date(year, MAY, mom_day)] = str_day

        # Children's Day
        if self.observed and date(year, SEP, 10):
            name = "Día del niño [Children day] (Observed)"
            self[date(year, SEP, 10)] = name

        # Independence Day
        if self.observed and date(year, SEP, 15):
            name = "Día de la Independencia [Independence Day]"
            self[date(year, SEP, 15)] = name

        # Teacher's Day
        if self.observed and date(year, SEP, 17):
            name = "Día del Maestro [Teacher's day] (Observed)"
            self[date(year, SEP, 17)] = name

        # October Holidays are joined on 3 days starting at October 3 to 6.
        # Some companies work medium day and take the rest on saturday.
        # This holiday is variant and some companies work normally.
        # If start day is weekend is ignored.
        # The main objective of this is to increase the tourism.

        # https://www.hondurastips.hn/2017/09/20/de-donde-nace-el-feriado-morazanico/

        if year <= 2014:
            # Morazan's Day
            if self.observed and date(year, OCT, 3):
                self[date(year, OCT, 3)] = "Día de Morazán [Morazan's Day]"

            # Columbus Day
            if self.observed and date(year, OCT, 12):
                self[date(year, OCT, 12)] = "Día de la Raza [Columbus Day]"

            # Amy Day
            if self.observed and date(year, OCT, 21):
                str_day = "Día de las Fuerzas Armadas [Army Day]"
                self[date(year, OCT, 21)] = str_day
        else:
            # Morazan Weekend
            if self.observed and date(year, OCT, 3):
                name = "Semana Morazánica [Morazan Weekend]"
                self[date(year, OCT, 3)] = name

            # Morazan Weekend
            if self.observed and date(year, OCT, 4):
                name = "Semana Morazánica [Morazan Weekend]"
                self[date(year, OCT, 4)] = name

            # Morazan Weekend
            if self.observed and date(year, OCT, 5):
                name = "Semana Morazánica [Morazan Weekend]"
                self[date(year, OCT, 5)] = name

        # Christmas
        self[date(year, DEC, 25)] = "Navidad [Christmas]"

class HN(Honduras):
    pass

class HND(Honduras):
    pass


# EDITOR > File and Code Templates put this here!

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

x=np.linspace(-10,10,100)
y=x**2

#%%
plt.close('all')
plt.ion()
plt.style.use('seaborn')  # style sheet reference: https://matplotlib.org/3.1.0/gallery/style_sheets/style_sheets_reference.html
fig, ax = plt.subplots( 1,1)#, figsize=(12,9), sharex=False)  #W x H
#fig = matplotlib.pyplot.figure

plt.plot(x,y, color='orange', linestyle='-')

# Do this, then get autocomplete pop-ups on plt
assert isinstance(ax, matplotlib.axes._subplots.Subplot)

# ax.

#%%


from __future__ import annotations

import configparser
import datetime
import os
import typing

import fabric
import invocations.console
import invoke

import invoke_patch
import src.constants

class GlobalConfig:
    host: str
    user: str
    key_filename: str
    project_name: str
    project_path: str

    source_filenames = [
        'main.py',
        'database.py',
        'utils.py',
        'telegram_utils.py',
        'analytics.py',
        'constants.py',

        'custom_logger.py',

        'config.cfg'
    ]
    meta_filenames = [
        'pyproject.toml',
        'poetry.lock'
    ]
    source_directories = [
        'migrations'
    ]

    @classmethod
    def load(cls) -> None:
        try:
            fabfile_config = configparser.ConfigParser()

            fabfile_config.read('fabfile.cfg')

            cls.host = fabfile_config.get('Fabric', 'Host')
            cls.user = fabfile_config.get('Fabric', 'User')
            cls.key_filename = os.path.expanduser(fabfile_config.get('Fabric', 'KeyFilename'))
            cls.project_name = fabfile_config.get('Fabric', 'ProjectName')
            cls.project_path = fabfile_config.get('Fabric', 'ProjectPath')
        except configparser.Error as error:
            raise invoke.Exit(
                message=f'Config error: {error}',
                code=1
            )

invoke_patch.fix_annotations()
GlobalConfig.load()

@fabric.task
def configure(connection: fabric.Connection) -> None:
    connection.user = GlobalConfig.user
    connection.inline_ssh_env = True
    connection.connect_kwargs.key_filename = GlobalConfig.key_filename

@fabric.task(pre=[configure], hosts=[GlobalConfig.host], help={'command': 'The shell command to execute on the server', 'env': 'An optional dictionary with environment variables'})
def execute(connection: fabric.Connection, command: str, env: typing.Dict[str, str] = None) -> None:
    if not command:
        return

    connection.run(command, env=env)

@fabric.task(pre=[configure], hosts=[GlobalConfig.host])
def cleanup(connection: fabric.Connection) -> None:
    question = f'Are you sure you want to completely delete the project "{GlobalConfig.project_name}" from "{GlobalConfig.host}"?'

    if invocations.console.confirm(
        question=question,
        assume_yes=False
    ):
        execute(connection, f'rm -rf {GlobalConfig.project_name}')
        execute(connection, f'rm -rf {GlobalConfig.project_path}/{GlobalConfig.project_name}')

@fabric.task(pre=[configure, cleanup], hosts=[GlobalConfig.host])
def setup(connection: fabric.Connection) -> None:
    execute(connection, f'mkdir -p {GlobalConfig.project_path}/{GlobalConfig.project_name}')
    execute(connection, f'ln -s {GlobalConfig.project_path}/{GlobalConfig.project_name} {GlobalConfig.project_name}')

    execute(connection, 'curl -sSL https://install.python-poetry.org | python -')

@fabric.task(pre=[configure], hosts=[GlobalConfig.host], help={'filename': 'An optional filename to deploy to the server'})
def upload(connection: fabric.Connection, filename: typing.Optional[str] = None) -> None:
    def upload_file(file_format: str, file_name: str, destination_path_format='{.project_name}/{}') -> None:
        connection.put(file_format.format(file_name), destination_path_format.format(GlobalConfig, file_name))

    def upload_directory(directory_name: str) -> None:
        execute(connection, f'mkdir -p {GlobalConfig.project_name}/{directory_name}')

        for _root, _directories, files in os.walk(f'src/{directory_name}'):
            for file in files:
                upload_file(f'src/{directory_name}/{{}}', file, f'{{.project_name}}/{directory_name}/{{}}')

    if filename:
        if filename in GlobalConfig.source_directories:
            upload_directory(filename)
        else:
            if filename in GlobalConfig.source_filenames:
                file_path_format = 'src/{}'
            elif filename in GlobalConfig.meta_filenames:
                file_path_format = '{}'
            else:
                raise invoke.ParseError(f'Filename "{filename}" is not registered')

            upload_file(file_path_format, filename)
    else:
        for name in GlobalConfig.source_filenames:
            upload_file('src/{}', name)

        for name in GlobalConfig.meta_filenames:
            upload_file('{}', name)

        for directory in GlobalConfig.source_directories:
            upload_directory(directory)

@fabric.task(pre=[configure], hosts=[GlobalConfig.host], help={'filename': 'An optional filename to deploy to the server'})
def deploy(connection: fabric.Connection, filename: typing.Optional[str] = None) -> None:
    upload(connection, filename)

    with connection.cd(GlobalConfig.project_name):
        execute(connection, 'eval "$(pyenv init --path)" && poetry install --no-dev', {
            'PATH': '$HOME/.pyenv/bin:$HOME/.poetry/bin:$PATH'
        })

@fabric.task(pre=[configure], hosts=[GlobalConfig.host], help={'filename': 'The filename to backup locally from the server'})
def backup(connection: fabric.Connection, filename: str) -> None:
    current_date = datetime.datetime.now().strftime(src.constants.GENERIC_DATE_FORMAT)
    name, extension = os.path.splitext(filename)

    with connection.cd(GlobalConfig.project_name):
        connection.get(f'{GlobalConfig.project_name}/{filename}', f'backup_{name}_{current_date}{extension}')

@fabric.task(pre=[configure], hosts=[GlobalConfig.host])
def backup_db(context: fabric.Connection) -> None:
    backup(context, 'file_convert.sqlite')


from django.urls import path
from . import views

app_name = 'quote'
urlpatterns = [
    path('', views.quote_page,),
    path('new', views.quote_page, name="new_quote"),
    path('schedule', views.schedule, name="schedule"),
    path('scheduling', views.scheduling, name="scheduling"),
    path('address', views.address, name="address"),
    path('setaddress', views.setaddress, name="setaddress"),
    path('revieworder', views.revieworder,  name="revieworder"),
    path('submitorder', views.submitorder,  name="submitorder"),
    path('receipt', views.receipt,  name="receipt"),
    path('billing', views.billing, name="billing"),

    path('clear', views.clear_quote,  name="clear_quote"),
    path('save', views.save,  name="save_quote"),
    path('saving', views.saving,  name="saving_quote"),
    path('savedquotes', views.savedquotes,  name="savedquotes"),
    path('delete/<int:quoteID>', views.destroyQuote,  name="destroyQuote"),
    path('edit/<int:quoteID>', views.editquote, name="edit_quote"),

    path('myaccount', views.account,  name="account"),
    path('manage_address', views.manage_address,  name="manage_address"),

    # AJAX to select a service and put it in session and then to table
    path('pickitem/<int:itemID>', views.pick_item,  name="pickitem"),
    path('updateitem/<int:itemID>/remove', views.remove_item,  name="removeitem"),
    path('updateitem/<int:itemID>', views.update_item,  name="updateitem"),
    path('updateitem/package/<str:package>/<int:itemID>', views.update_item_package,  name="updateitempackage"),
    path('updatequotetable', views.update_quote_table,  name="updateQuoteTable")
]


from os import environ as env
from tempfile import NamedTemporaryFile

from .papermill import execute
from . import gsmo

from utz import cd, sh

class Modules:
    def __init__(self, run=None, skip=None, conf=None):
        if isinstance(run, str): run = run.split(',')
        if isinstance(skip, str): skip = skip.split(',')
        if run is not None and skip is not None:
            raise RuntimeError('Specify at most one of {run,skip}: (%s, %s)' % (run, skip))
        self.runs = run
        self.skips = skip
        self.conf = conf or {}

    def run(self, module, nb='run.ipynb', out='nbs', dind=None, *args, **kwargs):
        if self.skips and module in self.skips:
            print(f'Module {module} marked as "skip"; skipping')
            return
        if self.runs and module not in self.runs:
            print(f'Module {module} not marked as "run"; skipping')
            return

        module_kwargs = self.conf.get(module, {})
        module_kwargs.update(kwargs)

        with cd(module):
            print(f'Running module: {module}')
            if dind is not False:
                with NamedTemporaryFile() as tmp:
                    with open(tmp.name,'w') as f:
                        import yaml
                        yaml.safe_dump(kwargs, f, sort_keys=False)
                    cmd = []
                    if 'GSMO_IMAGE' in env:
                        cmd += ['-i',env['GSMO_IMAGE']]
                    cmd += ['-I','run','-o',out,'-x',nb,'-Y',tmp.name]
                    gsmo.main(*cmd)
            else:
                execute(
                    nb,
                    out,
                    *args,
                    **module_kwargs,
                )
        sh('git','add',module)
        sh('git','commit','-m',module)

    def __call__(self, *args, **kwargs): return self.run(*args, **kwargs)


import torch.nn as nn
import torch.multiprocessing as mp
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler

import argparse
from lib.utils.file import bool_flag
from lib.utils.distributed import init_dist_node, init_dist_gpu, get_shared_folder
from lib.utils.flatten import flatten_dict

import submitit, random, sys
from pathlib import Path

def parse_args():
    parser = argparse.ArgumentParser(description='Template')
    parser.add_argument('--cfg', type=str, help='Configuration file')

    # PATHS
    paths_group = parser.add_argument_group("PATHS")
    paths_group.add_argument('--data', type=str, help='Path to raw dataset directory')
    paths_group.add_argument('--processed_data', type=str, help='Path to processed dataset directory')
    paths_group.add_argument('--processed_pattern', type=str, help='Pattern for processed data files')
    paths_group.add_argument('--genuine_pattern', type=str, help='Pattern for genuine data files')
    paths_group.add_argument('--simulated_pattern', type=str, help='Pattern for simulated data files')
    paths_group.add_argument('--out', type=str, help='Path to out directory')

    # GENERAL
    general_group = parser.add_argument_group("GENERAL")
    general_group.add_argument('--model', type=str, help='Model name')
    general_group.add_argument('--reset', action='store_true', help='Reset saved model logs and weights')
    general_group.add_argument('--tb', action='store_true', help='Start TensorBoard')
    general_group.add_argument('--gpus', type=str, help='GPUs list, only works if not on slurm')
    general_group.add_argument('--k_neighbors', type=int, help='Number of neighbors to consider in k-NN algorithm')

    # Model
    model_group = parser.add_argument_group("Model")
    model_group.add_argument('--num_layers', type=int, help='Number of layers')
    model_group.add_argument('--hidden_dim', type=int, help='Number of nodes in the hidden layer.')

    # Dataset
    dataset_group = parser.add_argument_group("Dataset")
    dataset_group.add_argument('--dataset', type=str, help='Dataset to choose')
    dataset_group.add_argument('--batch_per_gpu', type=int, help='Batch size per gpu')
    dataset_group.add_argument('--shuffle', action='store_true', help='Shuffle dataset')
    dataset_group.add_argument('--workers', type=int, help='Number of workers')
    dataset_group.add_argument('--keys', type=str, nargs='+', help='Columns to be used as data features')

    # Architecture
    architecture_group = parser.add_argument_group("Architecture")
    architecture_group.add_argument('--arch', type=str, help='Architecture to choose')

    # Trainer
    trainer_group = parser.add_argument_group("Trainer")
    trainer_group.add_argument('--trainer', type=str, help='Trainer to choose')
    trainer_group.add_argument('--epochs', type=int, help='Number of epochs')
    trainer_group.add_argument('--save_every', type=int, help='Save model every n epochs')
    trainer_group.add_argument('--fp16', action='store_true', help='Use fp16')

    # Optimization
    optimization_group = parser.add_argument_group("Optimization")
    optimization_group.add_argument('--optimizer', type=str, help='Optimizer to choose between "adam", "sgd", and "adagrad"')
    optimization_group.add_argument('--lr_start', type=float, help='Learning rate start')
    optimization_group.add_argument('--lr_end', type=float, help='Learning rate end')
    optimization_group.add_argument('--lr_warmup', type=int, help='Learning rate warmup')

    # SLURM
    slurm_group = parser.add_argument_group("SLURM")
    slurm_group.add_argument('--slurm', action='store_true', help='Use slurm')
    slurm_group.add_argument('--slurm_ngpus', type=int, help='Number of gpus per node')
    slurm_group.add_argument('--slurm_nnodes', type=int, help='Number of nodes')
    slurm_group.add_argument('--slurm_nodelist', type=str, help='Node list')
    slurm_group.add_argument('--slurm_partition', type=str, help='Partition')
    slurm_group.add_argument('--slurm_timeout', type=int, help='Timeout')

    args = parser.parse_args()

    # === Read CFG File === #
    if args.cfg:
        with open(args.cfg, 'r') as f:
            import ruamel.yaml as yaml
            yml = yaml.safe_load(f)

        # update values from cfg file only if not passed in cmdline
        cmd = [c[1:] for c in sys.argv if c[0]=='-']
        for k,v in flatten_dict(yml):
            if k not in cmd:
                args.__dict__[k] = v

    return args

class SLURM_Trainer(object):
    def __init__(self, args):
        self.args = args

    def __call__(self):

        init_dist_node(self.args)
        train(None, self.args)

def main():

    args = parse_args()
    args.port = random.randint(49152,65535)

    if args.slurm:

        # Almost copy-paste from https://github.com/facebookresearch/deit/blob/main/run_with_submitit.py
        args.output_dir = get_shared_folder(args) / "%j"
        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
        executor = submitit.AutoExecutor(folder=args.output_dir, slurm_max_num_timeout=30)

        if hasattr(args, "slurm_account"):
            executor.update_parameters(
                mem_gb=12*args.slurm_ngpus,
                gpus_per_node=args.slurm_ngpus,
                tasks_per_node=args.slurm_ngpus,
                cpus_per_task=2,
                nodes=args.slurm_nnodes,
                timeout_min=2800,
                slurm_partition=args.slurm_partition,
                account=args.slurm_account
            )
        else:
            executor.update_parameters(
                mem_gb=12*args.slurm_ngpus,
                gpus_per_node=args.slurm_ngpus,
                tasks_per_node=args.slurm_ngpus,
                cpus_per_task=2,
                nodes=args.slurm_nnodes,
                timeout_min=2800,
                slurm_partition=args.slurm_partition
            )

        if args.slurm_nodelist:
            executor.update_parameters(slurm_additional_parameters = {"nodelist": f'{args.slurm_nodelist}' })

        executor.update_parameters(name=args.model)
        trainer = SLURM_Trainer(args)
        job = executor.submit(trainer)
        print(f"Submitted job_id: {job.job_id}")

    else:
        init_dist_node(args)
        mp.spawn(train, args = (args,), nprocs = args.ngpus_per_node)

def train(gpu, args):

    # === SET ENV === #
    init_dist_gpu(gpu, args)

    # === DATA === #
    get_dataset = getattr(__import__("lib.datasets.{}".format(args.dataset), fromlist=["get_dataset"]), "get_dataset")
    dataset = get_dataset(args)

    sampler = DistributedSampler(dataset, shuffle=args.shuffle, num_replicas = args.world_size, rank = args.rank, seed = 31)
    loader = DataLoader(dataset=dataset,
                        sampler = sampler,
                        batch_size=args.batch_per_gpu,
                        num_workers= args.workers,
                        pin_memory = True,
                        drop_last = True
                       )
    print(f"Data loaded")

    # === MODEL === #
    get_model = getattr(__import__("lib.arch.{}".format(args.arch), fromlist=["get_model"]), "get_model")
    setattr(args, "input_dim", dataset.num_node_features)
    setattr(args, "outpup_dim", dataset.num_classes)
    setattr(args, "activation_function", nn.functional.relu)
    model = get_model(args).cuda(args.gpu)
    model = nn.SyncBatchNorm.convert_sync_batchnorm(model) # use if model contains batchnorm.
    model = nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])

    # === LOSS === #
    setattr(args, "loader", loader)
    from lib.core.loss import get_loss
    loss = get_loss(args).cuda(args.gpu)

    # === OPTIMIZER === #
    from lib.core.optimizer import get_optimizer
    optimizer = get_optimizer(model, args)

    # === TRAINING === #
    Trainer = getattr(__import__("lib.trainers.{}".format(args.trainer), fromlist=["Trainer"]), "Trainer")
    Trainer(args, loader, model, loss, optimizer).fit()

if __name__ == "__main__":
    main()


from typing import Optional, List
import numpy as np
import torch
import torch.nn as nn
from transformers import RagSequenceForGeneration, AutoModel, AutoConfig, RagRetriever, BatchEncoding
from transformers.models.dpr.modeling_dpr import DPRQuestionEncoderOutput
from distributed_pytorch_retriever import RagPyTorchDistributedRetriever

class RobertaRetriever(nn.Module):
  def __init__(self, model_name):
    super().__init__()

    config = AutoConfig.from_pretrained(model_name)
    self.encoder = AutoModel.from_pretrained(model_name)
    self.project = nn.Sequential(nn.Linear(config.hidden_size, config.hidden_size),
                                 nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps))

  def encode_seq(self, input_ids, mask):
    cls_rep = self.encoder(input_ids, mask)[0][:, 0, :]
    vector = self.project(cls_rep)
    return vector

  def forward(self,
              input_ids: Optional[torch.Tensor] = None,
              attention_mask: Optional[torch.Tensor] = None,
              token_type_ids: Optional[torch.Tensor] = None,
              inputs_embeds: Optional[torch.Tensor] = None,
              output_attentions=None,
              output_hidden_states=None,
              return_dict=None,):
    pooled = self.encode_seq(input_ids, attention_mask)
    if not return_dict:
      return [pooled]
    return DPRQuestionEncoderOutput(pooler_output=pooled, hidden_states=None, attentions=None)  # TODO: replace none

def load_saved(model, path, exact=True):
  try:
    state_dict = torch.load(path)
  except:
    state_dict = torch.load(path, map_location=torch.device('cpu'))

  def filter(x):
    if x.startswith('rag.question_encoder.'):
      return x[21:]
    if x.startswith('module.'):
      return x[7:]
    return x

  if exact:
    state_dict = {filter(k): v for (k, v) in state_dict.items()}
  else:
    state_dict = {filter(k): v for (k, v) in state_dict.items() if filter(k) in model.state_dict()}
  model.load_state_dict(state_dict, strict=False)  # TODO: embeddings.position_ids missing
  return model

class MyRagRetriever(RagRetriever):
  def __call__(
    self,
    question_input_ids: List[List[int]],
    question_hidden_states: np.ndarray,
    question_strings: List[str]=None,
    prefix=None,
    n_docs=None,
    return_tensors=None,
  ) -> BatchEncoding:
    n_docs = n_docs if n_docs is not None else self.n_docs
    prefix = prefix if prefix is not None else self.config.generator.prefix
    retrieved_doc_embeds, doc_ids, docs = self.retrieve(question_hidden_states, n_docs)

    if question_strings is None:
      input_strings = self.question_encoder_tokenizer.batch_decode(question_input_ids, skip_special_tokens=True)
    else:
      input_strings = [s.replace('[MASK]', '<mask>') for s in question_strings]  # TODO: make it more robust to different models
    context_input_ids, context_attention_mask = self.postprocess_docs(
      docs, input_strings, prefix, n_docs, return_tensors=return_tensors
    )

    return BatchEncoding(
      {
        "context_input_ids": context_input_ids,
        "context_attention_mask": context_attention_mask,
        "retrieved_doc_embeds": retrieved_doc_embeds,
        "doc_ids": doc_ids,
      },
      tensor_type=return_tensors,
    )

class MyRagPyTorchDistributedRetriever(RagPyTorchDistributedRetriever):
  def __call__(
    self,
    question_input_ids: List[List[int]],
    question_hidden_states: np.ndarray,
    question_strings: List[str]=None,
    prefix=None,
    n_docs=None,
    return_tensors=None,
  ) -> BatchEncoding:
    n_docs = n_docs if n_docs is not None else self.n_docs
    prefix = prefix if prefix is not None else self.config.generator.prefix
    retrieved_doc_embeds, doc_ids, docs = self.retrieve(question_hidden_states, n_docs)

    if question_strings is None:
      input_strings = self.question_encoder_tokenizer.batch_decode(question_input_ids, skip_special_tokens=True)
    else:
      input_strings = [s.replace('[MASK]', '<mask>') for s in question_strings]  # TODO: make it more robust to different models
    context_input_ids, context_attention_mask = self.postprocess_docs(
      docs, input_strings, prefix, n_docs, return_tensors=return_tensors
    )

    return BatchEncoding(
      {
        "context_input_ids": context_input_ids,
        "context_attention_mask": context_attention_mask,
        "retrieved_doc_embeds": retrieved_doc_embeds,
        "doc_ids": doc_ids,
      },
      tensor_type=return_tensors,
    )

class MemoryBank(nn.Module):
  def __init__(self, bank_size: int, emb_size: int = None):
    super().__init__()
    self.bank_size = bank_size
    self.emb_size = emb_size
    if emb_size is not None:
      self.register_buffer('queue', torch.randn(bank_size, emb_size))
    self.register_buffer('ptr', torch.zeros(1, dtype=torch.long))
    self.register_buffer('is_full', torch.zeros(1, dtype=bool))

  @property
  def is_initialized(self):
    return hasattr(self, 'queue')

  @torch.no_grad()
  def put(self,
          embeddings: torch.FloatTensor):  # (bs, emb_size)
    bs, es = embeddings.size()
    ptr = int(self.ptr)

    if not self.is_initialized:
      self.emb_size = es
      self.register_buffer('queue', torch.randn(self.bank_size, self.emb_size).to(embeddings.device))

    # update queue
    if ptr + bs > self.bank_size:  # reach the end of the queue
      self.queue[ptr:, :] = embeddings[:self.bank_size - ptr]
      self.queue[:ptr + bs - self.bank_size, :] = embeddings[self.bank_size - ptr:]
      self.is_full = self.is_full | True
    else:
      self.queue[ptr:ptr + bs, :] = embeddings

    # move pointer
    self.ptr[0] = (ptr + bs) % self.bank_size
    return

  @torch.no_grad()
  def get(self):
    if not self.is_initialized:
      raise ValueError('queue not initialized')
    if not self.is_full:
      return self.queue[:self.ptr]  # (<bank_size, emb_size)
    return self.queue

class MyRagSequenceForGeneration(RagSequenceForGeneration):
  def set_additional_retrievers(self, retrievers: List):
    self._retrievers = retrievers

  def retrieve_from_multiple(self, *args, **kwargs):
    if not hasattr(self, '_retrievers'):
      return self.retriever(*args, **kwargs)

    all_retrievers = [self.retriever] + self._retrievers
    question_hidden_states = torch.tensor(args[1])
    n_docs = kwargs['n_docs']

    cii_li = []
    cam_li = []
    rde_li = []
    di_li = []
    ds_li = []

    # collect retrieved results from all retrievers
    for retriever in all_retrievers:
      retriever_outputs = retriever(*args, **kwargs)
      cii, cam, rde, di = (
        retriever_outputs['context_input_ids'],  # (batch_size * n_docs, seq_len)
        retriever_outputs['context_attention_mask'],  # (batch_size * n_docs, seq_len)
        retriever_outputs['retrieved_doc_embeds'],  # (batch_size, n_docs, emb_size)
        retriever_outputs['doc_ids']  # (batch_size, n_docs)
      )
      cii_li.append(cii.view(-1, n_docs, cii.size(-1)))
      cam_li.append(cam.view(-1, n_docs, cam.size(-1)))
      rde_li.append(rde)
      di_li.append(di)
      doc_scores = torch.bmm(question_hidden_states.unsqueeze(1), rde.to(question_hidden_states).transpose(1, 2)).squeeze(1)
      ds_li.append(doc_scores)

    # merge results
    cii = torch.cat(cii_li, 1)
    cam = torch.cat(cam_li, 1)
    rde = torch.cat(rde_li, 1)
    di = torch.cat(di_li, 1)
    doc_scores = torch.cat(ds_li, 1)

    # sort by scores and select
    doc_scores, topk_ind = torch.topk(doc_scores, n_docs, dim=1)
    cii = torch.gather(cii, 1, topk_ind.unsqueeze(-1).repeat(1, 1, cii.size(-1))).view(-1, cii.size(-1))
    cam = torch.gather(cam, 1, topk_ind.unsqueeze(-1).repeat(1, 1, cam.size(-1))).view(-1, cam.size(-1))
    rde = torch.gather(rde, 1, topk_ind.unsqueeze(-1).repeat(1, 1, rde.size(-1)))
    di = torch.gather(di, 1, topk_ind)

    return {'context_input_ids': cii, 'context_attention_mask': cam, 'retrieved_doc_embeds': rde, 'doc_ids': di}

  @torch.no_grad()
  def generate(
    self,
    input_ids: Optional[torch.LongTensor] = None,
    attention_mask: Optional[torch.LongTensor] = None,
    context_input_ids=None,
    context_attention_mask=None,
    doc_scores=None,
    do_deduplication=None,  # defaults to True
    num_return_sequences=None,  # defaults to 1
    num_beams=None,  # defaults to 1
    n_docs=None,
    **model_kwargs
  ):
    """
    Implements RAG sequence "thorough" decoding. Read the :meth:`~transformers.PreTrainedModel.generate``
    documentation for more information on how to set other generate input parameters.
    Args:
        input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
            The sequence used as a prompt for the generation. If :obj:`input_ids` is not passed, then
            :obj:`context_input_ids` has to be provided.
        attention_mask (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
            Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
            - 1 for tokens that are **not masked**,
            - 0 for tokens that are **masked**.
            `What are attention masks? <../glossary.html#attention-mask>`__
        context_input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size * config.n_docs, config.max_combined_length)`, `optional`, returned when `output_retrieved=True`):
            Input IDs post-processed from the retrieved documents and the question encoder input_ids by the
            retriever.
        context_attention_mask (:obj:`torch.LongTensor` of shape :obj:`(batch_size * config.n_docs, config.max_combined_length)`, `optional`, returned when `output_retrieved=True`):
            Attention mask post-processed from the retrieved documents and the question encoder :obj:`input_ids` by
            the retriever.
            If the model is not initialized with a ``retriever`` or ``input_ids`` is not given,
            :obj:`context_input_ids` and :obj:`context_attention_mask` have to be provided to the forward pass.
            They are returned by :meth:`~transformers.RagRetriever.__call__`.
        doc_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, config.n_docs)`):
            Score between each retrieved document embeddings (see :obj:`retrieved_doc_embeds`) and
            :obj:`question_encoder_last_hidden_state`.
            If the model is not initialized with a ``retriever`` or ``input_ids`` is not given, :obj:`doc_scores`
            has to be provided to the forward pass. :obj:`doc_scores` are returned by
            :meth:`~transformers.RagRetriever.__call__`.
        do_deduplication (:obj:`bool`, `optional`):
            Whether or not to deduplicate the generations from different context documents for a given input. Has
            to be set to :obj:`False` if used while training with distributed backend.
        num_return_sequences(:obj:`int`, `optional`, defaults to 1):
            The number of independently computed returned sequences for each element in the batch. Note that this
            is not the value we pass to the ``generator``'s `:func:`~transformers.PreTrainedModel.generate``
            function, where we set ``num_return_sequences`` to :obj:`num_beams`.
        num_beams (:obj:`int`, `optional`, defaults to 1):
            Number of beams for beam search. 1 means no beam search.
        n_docs (:obj:`int`, `optional`, defaults to :obj:`config.n_docs`)
            Number of documents to retrieve and/or number of documents for which to generate an answer.
        kwargs:
            Additional kwargs will be passed to :meth:`~transformers.PreTrainedModel.generate`.
    Return:
        :obj:`torch.LongTensor` of shape :obj:`(batch_size * num_return_sequences, sequence_length)`: The generated
        sequences. The second dimension (sequence length) is either equal to :obj:`max_length` or shorter if all
        batches finished early due to the :obj:`eos_token_id`.
    """

    n_docs = n_docs if n_docs is not None else self.config.n_docs
    do_deduplication = do_deduplication if do_deduplication is not None else self.config.do_deduplication
    num_doc_return_sequences = (
      num_return_sequences if num_return_sequences is not None else self.config.num_return_sequences
    )
    num_beams = num_beams if num_beams is not None else self.config.num_beams

    assert (
      input_ids is not None or context_input_ids is not None
    ), " At least one of input_ids or context_input_ids must be given"

    if self.retriever is not None and context_input_ids is None:
      question_hidden_states = self.question_encoder(input_ids, attention_mask=attention_mask)[0]
      context_input_ids = self.retriever(
        input_ids,
        question_hidden_states.cpu().detach().to(torch.float32).numpy(),
        prefix=self.generator.config.prefix,
        n_docs=n_docs,
        return_tensors="pt",
      )["context_input_ids"]

      # set to correct device
      context_input_ids = context_input_ids.to(input_ids)

    hypos = []
    logprobs = []
    model_kwargs["num_beams"] = num_beams
    model_kwargs["num_return_sequences"] = num_beams
    model_kwargs["attention_mask"] = None

    batch_size = input_ids.shape[0] if input_ids is not None else context_input_ids.shape[0] // n_docs

    for index in range(batch_size):
      # first, generate beams from documents:
      generator_input_ids = context_input_ids[index * n_docs: (index + 1) * n_docs]  # (n_docs, max_len)

      output_sequences = self.generator.generate(
        generator_input_ids,
        **model_kwargs,
      )  # n_docs * n_beam, tgt_len
      if do_deduplication:
        # do_deduplication, max_output_len
        output_sequences = torch.stack(list({str(k.tolist()): k for k in output_sequences}.values()))

      num_candidates = output_sequences.shape[
        0
      ]  # after deduplication, this number can be less than n_docs*n_beam

      # then, run model forwards to get nll scores:
      if input_ids is not None:
        new_input_ids = input_ids[index: index + 1].repeat(num_candidates, 1)
        outputs = self(new_input_ids, labels=output_sequences, exclude_bos_score=True)
      else:  # input_ids is None, need context_input_ids/mask and doc_scores
        assert (
          context_attention_mask is not None
        ), "Make sure that `context_attention_mask` are passed, if no `input_ids` is set. Alternatively, you can set a retriever using the `set_retriever(...)` function."
        assert (
          doc_scores is not None
        ), "Make sure that `doc_scores` are passed, if no `input_ids` is set. Alternatively, you can set a retriever using the `set_retriever(...)` function."

        individual_input_ids = generator_input_ids.repeat(
          num_candidates, 1
        )  # (num_candidates*n_docs, max_len)

        individual_attention_mask = context_attention_mask[index * n_docs: (index + 1) * n_docs]
        individual_attention_mask = individual_attention_mask.repeat(num_candidates, 1)

        individual_doc_scores = doc_scores[index: (index + 1), :]  # doc_scores.shape = [batch, n_docs]
        individual_doc_scores = individual_doc_scores.repeat(num_candidates, 1)  # [num_candidates, n_docs]

        outputs = self(
          context_input_ids=individual_input_ids,
          context_attention_mask=individual_attention_mask,
          doc_scores=individual_doc_scores,
          labels=output_sequences,
          exclude_bos_score=True,
        )

      lps, top_cand_inds = (-outputs["loss"]).topk(num_doc_return_sequences)

      # add hypothesis
      hypos.append(output_sequences[top_cand_inds])
      logprobs.append(lps)

    logprobs = torch.cat(logprobs, 0)
    return self._cat_and_pad(hypos, pad_token_id=self.config.generator.pad_token_id), logprobs

  @staticmethod
  def load_question_encoder(rag_model, question_encoder_path: str):
    model = RobertaRetriever('roberta-base')
    model = load_saved(model, question_encoder_path, exact=False)
    rag_model.rag.question_encoder = model

  def get_nll(
    self, seq_logits, doc_scores, target, reduce_loss=False, epsilon=0.0, exclude_bos_score=False, n_docs=None
  ):
    # shift tokens left
    target = torch.cat(
      [target[:, 1:], target.new(target.shape[0], 1).fill_(self.config.generator.pad_token_id)], 1
    )

    n_docs = n_docs if n_docs is not None else self.config.n_docs

    # bos_token_id is None for T5
    bos_token_id = self.config.bos_token_id or self.config.generator.bos_token_id
    use_bos = bos_token_id is not None and target[:, 0].eq(bos_token_id).all()

    def _mask_pads(ll, smooth_obj):
      pad_mask = target.eq(self.config.generator.pad_token_id)
      if pad_mask.any():
        ll.masked_fill_(pad_mask, 0.0)
        smooth_obj.masked_fill_(pad_mask, 0.0)
      return ll.squeeze(-1), smooth_obj.squeeze(-1)

    # seq_logits dim = (batch*n_docs, tgt_len , #vocabs)
    seq_logprobs = torch.nn.functional.log_softmax(seq_logits, dim=-1).view(
      seq_logits.shape[0] // n_docs, n_docs, -1, seq_logits.size(-1)
    )  # batch_size x n_docs x tgt_len x #vocab_size

    # (bs, n_docs) or (bs, bs * n_docs) or (bs, bs * n_docs + memory_bank_size)
    doc_logprobs = torch.nn.functional.log_softmax(doc_scores, dim=1)

    bs = doc_logprobs.size(0)
    if doc_logprobs.size(1) == n_docs:
      pass
    elif doc_logprobs.size(1) >= bs * n_docs:  # contain doc log probs over the whole batch (plus memory bank)
      # (bs, bs, n_docs)
      doc_logprobs = doc_logprobs[:, :bs * n_docs].view(bs, bs, n_docs)  # remove memory bank
      # (bs, n_docs)
      doc_logprobs = torch.masked_select(
        doc_logprobs,
        torch.eye(bs).unsqueeze(-1).bool().to(doc_logprobs.device)).view(bs, n_docs)
    else:
      raise Exception(f'the size of document log prob {doc_logprobs.size()} is unexpected')

    doc_logprobs = doc_logprobs.unsqueeze(-1).unsqueeze(-1)

    # RAG-sequence marginalization
    first_token_scores = seq_logprobs[:, :, :1, :]
    second_token_scores = seq_logprobs[:, :, 1:2, :]
    remainder = seq_logprobs[:, :, 2:, :]
    rag_logprobs = torch.cat([first_token_scores, second_token_scores + doc_logprobs, remainder], dim=2)

    # calculate loss
    target = target.unsqueeze(1).unsqueeze(-1).repeat(1, n_docs, 1, 1)
    assert target.dim() == rag_logprobs.dim()

    ll = rag_logprobs.gather(dim=-1, index=target)
    smooth_obj = rag_logprobs.sum(dim=-1, keepdim=True)  # total sum of all (normalised) logits

    ll, smooth_obj = _mask_pads(ll, smooth_obj)

    # sum over tokens, exclude bos while scoring
    ll = ll[:, :, 1:].sum(2) if exclude_bos_score and use_bos else ll.sum(2)
    smooth_obj = smooth_obj.sum(2)
    ll = ll.logsumexp(1)  # logsumexp over docs
    smooth_obj = smooth_obj.logsumexp(1)

    nll_loss = -ll
    smooth_loss = -smooth_obj

    if reduce_loss:
      nll_loss = nll_loss.sum()
      smooth_loss = smooth_loss.sum()

    eps_i = epsilon / rag_logprobs.size(-1)
    loss = (1.0 - epsilon) * nll_loss + eps_i * smooth_loss
    return loss


from zope.component import adapts
from zope.interface import implements

from Products.ZenModel.RRDDataSource import SimpleRRDDataSource
from Products.ZenModel.ZenPackPersistence import ZenPackPersistence
from Products.ZenUtils.ZenTales import talesEvalStr
from Products.Zuul.form import schema
from Products.Zuul.infos import ProxyProperty
from Products.Zuul.infos.template import RRDDataSourceInfo
from Products.Zuul.interfaces import IRRDDataSourceInfo
from Products.Zuul.utils import ZuulMessageFactory as _t
from Products.ZenWidgets import messaging

from Products.Zuul.interfaces import ICatalogTool, IInfo
from Products.AdvancedQuery import Eq

class MonitoredPropertyDataSource(ZenPackPersistence, SimpleRRDDataSource):
    '''
    Model class for MonitoredPropertyDataSource.
    '''

    ZENPACKID = 'ZenPacks.zenoss.PropertyMonitor'

    MBP_TYPE = 'Property'

    sourcetypes = (MBP_TYPE,)
    sourcetype = MBP_TYPE

    # SimpleRRDDataSource property overrides.
    cycletime = '${here/zPropertyMonitorInterval}'
    eventClass = '/Ignore'
    severity = 0

    class_name = ''
    property_name = ''

    _properties = tuple([p for p in SimpleRRDDataSource._properties if p['id'] != 'cycletime']) + (
        {'id': 'cycletime', 'type': 'string', 'mode': 'w'},
        {'id': 'class_name', 'type': 'string', 'mode': 'w'},
        {'id': 'property_name', 'type': 'string', 'mode': 'w'},
    )

    def getDescription(self):
        '''
        Return a friendly description of this datasource.
        '''
        if not self.property_name:
            return 'Not Configured'

        return "%s / %s" % (self.class_name, self.property_name)

    def getComponent(self, context):
        return context.id

    def talesEval(self, text, context):
        device = context.device()
        extra = {
            'device': device,
            'dev': device,
            'devname': device.id,
            'datasource': self,
            'ds': self,
            }

        return talesEvalStr(str(text), context, extra=extra)

    def getCycleTime(self, context):
        return int(self.talesEval(self.cycletime, context))

    def testDataSourceAgainstDevice(self, testDevice, REQUEST, write, errorLog):
        """
        Does the majority of the logic for testing a datasource against the device
        @param string testDevice The id of the device we are testing
        @param Dict REQUEST the browers request
        @param Function write The output method we are using to stream the result of the command
        @parma Function errorLog The output method we are using to report errors
        """
        out = REQUEST.RESPONSE
        # Determine which device to execute against
        device = None
        if testDevice:
            # Try to get specified device
            device = self.findDevice(testDevice)
            if not device:
                errorLog(
                    'No device found',
                    'Cannot find device matching %s.' % testDevice,
                    priority=messaging.WARNING
                )
                return self.callZenScreen(REQUEST)
        elif hasattr(self, 'device'):
            # ds defined on a device, use that device
            device = self.device()
        elif hasattr(self, 'getSubDevicesGen'):
            # ds defined on a device class, use any device from the class
            try:
                device = self.getSubDevicesGen().next()
            except StopIteration:
                # No devices in this class, bail out
                pass
        if not device:
            errorLog(
                'No Testable Device',
                'Cannot determine a device against which to test.',
                priority=messaging.WARNING
            )
            return self.callZenScreen(REQUEST)

        ## Execute the datasource

        class_name = REQUEST.get('class_name')
        property_name = REQUEST.get('property_name')
        results = ICatalogTool(device).search(query=Eq('meta_type', class_name))

        if not results.total:
            out.write("0 objects found. No test performed.\n")
            return

        if results.total > 10:
            out.write("%d %s components found, showing first 10<p>\n" % (results.total, class_name))

        for (i, result) in enumerate(results):
            obj = result.getObject()
            info = IInfo(obj)
            property_value = getattr(info, property_name, "ERROR")

            out.write("    '%s' %s = %s<br>\n" % (obj.titleOrId(), property_name, property_value))

            if i + 1 == 10:
                break

class IMonitoredPropertyDataSourceInfo(IRRDDataSourceInfo):
    '''
    API Info interface for MonitoredPropertyDataSource.
    '''

    # IRRDDataSourceInfo doesn't define this.
    cycletime = schema.TextLine(title=_t(u'Cycle Time (seconds)'))

    # The xtype for class_name also manages property_name.
    class_name = schema.TextLine(
        title=_t(u'Property'),
        group=_t('Detail'),
        xtype='mbp_property')

class MonitoredPropertyDataSourceInfo(RRDDataSourceInfo):
    '''
    API Info adapter factory for MonitoredPropertyDataSource.
    '''

    implements(IMonitoredPropertyDataSourceInfo)
    adapts(MonitoredPropertyDataSource)

    # RRDDataSourceInfo doesn't define this.
    cycletime = ProxyProperty('cycletime')

    class_name = ProxyProperty('class_name')
    property_name = ProxyProperty('property_name')

    @property
    def testable(self):
        """
        This tells the client if we can test this datasource against a
        specific device.
        """
        return True


from fastapi import Depends, FastAPI
from fastapi.middleware.cors import CORSMiddleware

from auth.verify_token import get_query_token

import routers.s3_objects_router as s3_objects_router
import routers.users_router as users_router
import routers.user_objects_router as user_objects_router

import uvicorn

PORT = 8000

app = FastAPI(dependencies=[Depends(get_query_token)])

# origins = [
#     "http://localhost",
#     "http://localhost:3000",
# ]

origins = ["*"]

app.add_middleware(
    CORSMiddleware,

    allow_origins=origins,
    allow_credentials=True,

    allow_methods=["*"],
    allow_headers=["*"],
)

app.include_router(s3_objects_router.router)
app.include_router(users_router.router)
app.include_router(user_objects_router.router)

@app.get('/')
async def root():
    return {'message': 'Welcome to Hyper Wasabi!'}

if __name__ == '__main__':
    uvicorn.run('main:app', port=PORT, reload=True)


import base64
import wget
import settings

weights_path = settings.WEIGHTS_PATH

def from_onedrive(weights_path):

    data_bytes64 = base64.b64encode(bytes(settings.ONEDRIVE_LINK, "utf-8"))
    data_bytes64_String = (
        data_bytes64.decode("utf-8").replace("/", "_").replace("+", "-").rstrip("=")
    )

    resultUrl = (
        f"https://api.onedrive.com/v1.0/shares/u!{data_bytes64_String}/root/content"
    )

    wget.download(resultUrl, weights_path)