<commit_before># stdlib
from typing import Any
from typing import Iterable
from typing import KeysView
from typing import List
from typing import Optional
from typing import Union
from typing import cast

# third party
import redis
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm.session import Session
from torch import Tensor

# syft absolute
import syft

# relative
from ....common.uid import UID
from ....node.common.node_table.bin_obj_dataset import BinObjDataset
from ....store import ObjectStore
from ....store.storeable_object import StorableObject
from ..node_table.bin_obj_metadata import ObjectMetadata

ENCODING = "UTF-8"

def create_storable(
    _id: UID, data: Tensor, description: str, tags: Optional[List[str]] = None
) -> StorableObject:
    obj = StorableObject(id=_id, data=data, description=description, tags=tags)

    return obj

class DictStore(ObjectStore):
    def __init__(self, db: Session) -> None:
        self.db = db
        self.kv_store: dict[UID, Any] = {}

    def get_object(self, key: UID) -> Optional[StorableObject]:
        try:
            return self.__getitem__(key)
        except KeyError as e:  # noqa: F841
            return None

    def get_objects_of_type(self, obj_type: type) -> Iterable[StorableObject]:
        # raise NotImplementedError("get_objects_of_type")
        # return [obj for obj in self.values() if isinstance(obj.data, obj_type)]
        return self.values()

    def __sizeof__(self) -> int:
        return self.values().__sizeof__()

    def __str__(self) -> str:
        print("DictStore __str__")

    def __len__(self) -> int:
        return len(self.kv_store.keys())

    def keys(self) -> KeysView[UID]:
        return self.kv_store.keys()

    def values(self) -> List[StorableObject]:
        return self.kv_store.values()

    def __contains__(self, key: UID) -> bool:
        return key in self.kv_store

    def __getitem__(self, key: Union[UID, str, bytes]) -> StorableObject:
        try:
            local_session = sessionmaker(bind=self.db)()
            # bin_obj = local_session.query(BinObject).filter_by(id=str(key.value)).first()

            if isinstance(key, UID):
                key_str = str(key.value)
                key_uid = key
            elif isinstance(key, bytes):
                key_str = str(key.decode("utf-8"))
                key_uid = UID.from_string(key_str)
            else:
                key_str = key
                key_uid = UID.from_string(key_str)

            obj = self.kv_store[key_uid]
            obj_metadata = (
                local_session.query(ObjectMetadata).filter_by(obj=key_str).first()
            )

            if obj is None or obj_metadata is None:
                raise KeyError(f"Object not found! for UID: {key_uid}")

            obj = StorableObject(
                id=key_uid,
                data=obj,
                description=obj_metadata.description,
                tags=obj_metadata.tags,
                read_permissions=dict(
                    syft.deserialize(
                        bytes.fromhex(obj_metadata.read_permissions), from_bytes=True
                    )
                ),
                search_permissions=dict(
                    syft.deserialize(
                        bytes.fromhex(obj_metadata.search_permissions), from_bytes=True
                    )
                ),
                # name=obj_metadata.name,
            )
            local_session.close()
            return obj
        except Exception as e:
            print(f"Cant store object {key}", e)
            raise KeyError(f"Object not found! for UID: {key}")

    def is_dataset(self, key: UID) -> bool:
        local_session = sessionmaker(bind=self.db)()
        is_dataset_obj = (
            local_session.query(BinObjDataset).filter_by(obj=str(key.value)).exists()
        )
        is_dataset_obj = local_session.query(is_dataset_obj).scalar()
        local_session.close()
        return is_dataset_obj

    def _get_obj_dataset_relation(self, key: UID) -> Optional[BinObjDataset]:
        local_session = sessionmaker(bind=self.db)()
        obj_dataset_relation = (
            local_session.query(BinObjDataset).filter_by(obj=str(key.value)).first()
        )
        local_session.close()
        return obj_dataset_relation

    def __setitem__(self, key: UID, value: StorableObject) -> None:
        try:
            self.kv_store[key] = value
        except Exception as e:
            print(f"failed to add {key} and value to dict", e)

        metadata_obj = ObjectMetadata(
            obj=str(key.value),
            tags=value.tags,
            description=value.description,
            read_permissions=cast(
                bytes,
                syft.serialize(
                    syft.lib.python.Dict(value.read_permissions), to_bytes=True
                ),
            ).hex(),
            search_permissions=cast(
                bytes,
                syft.serialize(
                    syft.lib.python.Dict(value.search_permissions), to_bytes=True
                ),
            ).hex(),
            # name=metadata_dict["name"],
        )

        obj_dataset_relation = self._get_obj_dataset_relation(key)
        if obj_dataset_relation:
            # Create a object dataset relationship for the new object
            obj_dataset_relation = BinObjDataset(
                id=obj_dataset_relation.id,
                name=obj_dataset_relation.name,
                obj=key,
                dataset=obj_dataset_relation.dataset,
                dtype=obj_dataset_relation.dtype,
                shape=obj_dataset_relation.shape,
            )

        local_session = sessionmaker(bind=self.db)()
        local_session.add(metadata_obj)
        local_session.add(obj_dataset_relation) if obj_dataset_relation else None
        local_session.commit()
        local_session.close()

    def delete(self, key: UID) -> None:
        try:
            del self.kv_store[key]
            local_session = sessionmaker(bind=self.db)()
            metadata_to_delete = (
                local_session.query(ObjectMetadata)
                .filter_by(obj=str(key.value))
                .first()
            )
            local_session.delete(metadata_to_delete)
            local_session.commit()
            local_session.close()
        except Exception as e:
            print(f"{type(self)} Exception in __delitem__ error {key}. {e}")

    def clear(self) -> None:
        self.kv_store = {}
        local_session = sessionmaker(bind=self.db)()
        local_session.query(ObjectMetadata).delete()
        local_session.commit()
        local_session.close()

    def __repr__(self) -> str:
        return "DictStore __repr__"

class RedisStore(ObjectStore):
    def __init__(self, db: Session) -> None:
        self.db = db
        print("connecting to redis")
        try:
            self.redis = redis.Redis(host="redis", port=6379)
        except Exception as e:
            print("failed to load redis", e)
            raise e

    def get_object(self, key: UID) -> Optional[StorableObject]:
        try:
            return self.__getitem__(key)
        except KeyError as e:  # noqa: F841
            return None

    def get_objects_of_type(self, obj_type: type) -> Iterable[StorableObject]:
        # raise NotImplementedError("get_objects_of_type")
        # return [obj for obj in self.values() if isinstance(obj.data, obj_type)]
        return self.values()

    def __sizeof__(self) -> int:
        return self.values().__sizeof__()

    def __str__(self) -> str:
        return str(self.values())

    def __len__(self) -> int:
        return self.client.dbsize()

    def keys(self) -> KeysView[UID]:
        key_bytes = self.redis.keys()
        key_ids = [UID.from_string(str(key.decode("utf-8"))) for key in key_bytes]
        return key_ids

    def values(self) -> List[StorableObject]:
        key_bytes = self.redis.keys()
        # this is bad we need to decouple getting the data from the search
        all_values = []
        for key in key_bytes:
            all_values.append(self.__getitem__(key))

        # return self.redis.mget(self.keys())
        return all_values

    def __contains__(self, key: UID) -> bool:
        return self.redis.contains(str(key.value))

    def __getitem__(self, key: Union[UID, str, bytes]) -> StorableObject:
        local_session = sessionmaker(bind=self.db)()

        if isinstance(key, UID):
            key_str = str(key.value)
            key_uid = key
        elif isinstance(key, bytes):
            key_str = str(key.decode("utf-8"))
            key_uid = UID.from_string(key_str)
        else:
            key_str = key
            key_uid = UID.from_string(key_str)

        bin = self.redis.get(key_str)
        obj = syft.deserialize(bin, from_bytes=True)
        obj_metadata = (
            local_session.query(ObjectMetadata).filter_by(obj=key_str).first()
        )

        if obj is None or obj_metadata is None:
            raise KeyError(f"Object not found! for UID: {key_uid}")

        obj = StorableObject(
            id=key_uid,
            data=obj,
            description=obj_metadata.description,
            tags=obj_metadata.tags,
            read_permissions=dict(
                syft.deserialize(
                    bytes.fromhex(obj_metadata.read_permissions), from_bytes=True
                )
            ),
            search_permissions=dict(
                syft.deserialize(
                    bytes.fromhex(obj_metadata.search_permissions), from_bytes=True
                )
            ),
            # name=obj_metadata.name,
        )
        local_session.close()
        return obj

    def is_dataset(self, key: UID) -> bool:
        local_session = sessionmaker(bind=self.db)()
        is_dataset_obj = (
            local_session.query(BinObjDataset).filter_by(obj=str(key.value)).exists()
        )
        is_dataset_obj = local_session.query(is_dataset_obj).scalar()
        local_session.close()
        return is_dataset_obj

    def _get_obj_dataset_relation(self, key: UID) -> Optional[BinObjDataset]:
        local_session = sessionmaker(bind=self.db)()
        obj_dataset_relation = (
            local_session.query(BinObjDataset).filter_by(obj=str(key.value)).first()
        )
        local_session.close()
        return obj_dataset_relation

    def __setitem__(self, key: UID, value: StorableObject) -> None:

        bin = syft.serialize(value.data, to_bytes=True)
        self.redis.set(str(key.value), bin)

        metadata_obj = ObjectMetadata(
            obj=str(key.value),
            tags=value.tags,
            description=value.description,
            read_permissions=cast(
                bytes,
                syft.serialize(
                    syft.lib.python.Dict(value.read_permissions), to_bytes=True
                ),
            ).hex(),
            search_permissions=cast(
                bytes,
                syft.serialize(
                    syft.lib.python.Dict(value.search_permissions), to_bytes=True
                ),
            ).hex(),
            # name=metadata_dict["name"],
        )

        obj_dataset_relation = self._get_obj_dataset_relation(key)
        if obj_dataset_relation:
            # Create a object dataset relationship for the new object
            obj_dataset_relation = BinObjDataset(
                id=obj_dataset_relation.id,
                name=obj_dataset_relation.name,
                obj=key,
                dataset=obj_dataset_relation.dataset,
                dtype=obj_dataset_relation.dtype,
                shape=obj_dataset_relation.shape,
            )

        local_session = sessionmaker(bind=self.db)()
        local_session.add(metadata_obj)
        local_session.add(obj_dataset_relation) if obj_dataset_relation else None
        local_session.commit()
        local_session.close()

    def delete(self, key: UID) -> None:
        try:
            self.redis.delete(str(key.value))
            local_session = sessionmaker(bind=self.db)()
            metadata_to_delete = (
                local_session.query(ObjectMetadata)
                .filter_by(obj=str(key.value))
                .first()
            )
            local_session.delete(metadata_to_delete)
            local_session.commit()
            local_session.close()
        except Exception as e:
            print(f"{type(self)} Exception in __delitem__ error {key}. {e}")

    def clear(self) -> None:
        self.redis.flushdb()
        local_session = sessionmaker(bind=self.db)()
        local_session.query(ObjectMetadata).delete()
        local_session.commit()
        local_session.close()

    def __repr__(self) -> str:
        return "RedisStore __repr__"
<commit_msg>Fixed __str__ and prevented loading all values by accident<commit_after># stdlib
from typing import Any
from typing import Iterable
from typing import KeysView
from typing import List
from typing import Optional
from typing import Union
from typing import cast

# third party
import redis
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm.session import Session
from torch import Tensor

# syft absolute
import syft

# relative
from ....common.uid import UID
from ....node.common.node_table.bin_obj_dataset import BinObjDataset
from ....store import ObjectStore
from ....store.storeable_object import StorableObject
from ..node_table.bin_obj_metadata import ObjectMetadata

ENCODING = "UTF-8"

def create_storable(
    _id: UID, data: Tensor, description: str, tags: Optional[List[str]] = None
) -> StorableObject:
    obj = StorableObject(id=_id, data=data, description=description, tags=tags)

    return obj

class DictStore(ObjectStore):
    def __init__(self, db: Session) -> None:
        self.db = db
        self.kv_store: dict[UID, Any] = {}

    def get_object(self, key: UID) -> Optional[StorableObject]:
        try:
            return self.__getitem__(key)
        except KeyError as e:  # noqa: F841
            return None

    def get_objects_of_type(self, obj_type: type) -> Iterable[StorableObject]:
        # raise NotImplementedError("get_objects_of_type")
        # return [obj for obj in self.values() if isinstance(obj.data, obj_type)]
        return self.values()

    def __sizeof__(self) -> int:
        return self.values().__sizeof__()

    def __str__(self) -> str:
        return f"{type(self)}"

    def __len__(self) -> int:
        return len(self.kv_store.keys())

    def keys(self) -> KeysView[UID]:
        return self.kv_store.keys()

    def values(self) -> List[StorableObject]:
        return self.kv_store.values()

    def __contains__(self, key: UID) -> bool:
        return key in self.kv_store

    def __getitem__(self, key: Union[UID, str, bytes]) -> StorableObject:
        try:
            local_session = sessionmaker(bind=self.db)()
            # bin_obj = local_session.query(BinObject).filter_by(id=str(key.value)).first()

            if isinstance(key, UID):
                key_str = str(key.value)
                key_uid = key
            elif isinstance(key, bytes):
                key_str = str(key.decode("utf-8"))
                key_uid = UID.from_string(key_str)
            else:
                key_str = key
                key_uid = UID.from_string(key_str)

            obj = self.kv_store[key_uid]
            obj_metadata = (
                local_session.query(ObjectMetadata).filter_by(obj=key_str).first()
            )

            if obj is None or obj_metadata is None:
                raise KeyError(f"Object not found! for UID: {key_uid}")

            obj = StorableObject(
                id=key_uid,
                data=obj,
                description=obj_metadata.description,
                tags=obj_metadata.tags,
                read_permissions=dict(
                    syft.deserialize(
                        bytes.fromhex(obj_metadata.read_permissions), from_bytes=True
                    )
                ),
                search_permissions=dict(
                    syft.deserialize(
                        bytes.fromhex(obj_metadata.search_permissions), from_bytes=True
                    )
                ),
                # name=obj_metadata.name,
            )
            local_session.close()
            return obj
        except Exception as e:
            print(f"Cant store object {key}", e)
            raise KeyError(f"Object not found! for UID: {key}")

    def is_dataset(self, key: UID) -> bool:
        local_session = sessionmaker(bind=self.db)()
        is_dataset_obj = (
            local_session.query(BinObjDataset).filter_by(obj=str(key.value)).exists()
        )
        is_dataset_obj = local_session.query(is_dataset_obj).scalar()
        local_session.close()
        return is_dataset_obj

    def _get_obj_dataset_relation(self, key: UID) -> Optional[BinObjDataset]:
        local_session = sessionmaker(bind=self.db)()
        obj_dataset_relation = (
            local_session.query(BinObjDataset).filter_by(obj=str(key.value)).first()
        )
        local_session.close()
        return obj_dataset_relation

    def __setitem__(self, key: UID, value: StorableObject) -> None:
        try:
            self.kv_store[key] = value
        except Exception as e:
            print(f"failed to add {key} and value to dict", e)

        metadata_obj = ObjectMetadata(
            obj=str(key.value),
            tags=value.tags,
            description=value.description,
            read_permissions=cast(
                bytes,
                syft.serialize(
                    syft.lib.python.Dict(value.read_permissions), to_bytes=True
                ),
            ).hex(),
            search_permissions=cast(
                bytes,
                syft.serialize(
                    syft.lib.python.Dict(value.search_permissions), to_bytes=True
                ),
            ).hex(),
            # name=metadata_dict["name"],
        )

        obj_dataset_relation = self._get_obj_dataset_relation(key)
        if obj_dataset_relation:
            # Create a object dataset relationship for the new object
            obj_dataset_relation = BinObjDataset(
                id=obj_dataset_relation.id,
                name=obj_dataset_relation.name,
                obj=key,
                dataset=obj_dataset_relation.dataset,
                dtype=obj_dataset_relation.dtype,
                shape=obj_dataset_relation.shape,
            )

        local_session = sessionmaker(bind=self.db)()
        local_session.add(metadata_obj)
        local_session.add(obj_dataset_relation) if obj_dataset_relation else None
        local_session.commit()
        local_session.close()

    def delete(self, key: UID) -> None:
        try:
            del self.kv_store[key]
            local_session = sessionmaker(bind=self.db)()
            metadata_to_delete = (
                local_session.query(ObjectMetadata)
                .filter_by(obj=str(key.value))
                .first()
            )
            local_session.delete(metadata_to_delete)
            local_session.commit()
            local_session.close()
        except Exception as e:
            print(f"{type(self)} Exception in __delitem__ error {key}. {e}")

    def clear(self) -> None:
        self.kv_store = {}
        local_session = sessionmaker(bind=self.db)()
        local_session.query(ObjectMetadata).delete()
        local_session.commit()
        local_session.close()

    def __repr__(self) -> str:
        return f"{type(self)}"

class RedisStore(ObjectStore):
    def __init__(self, db: Session) -> None:
        self.db = db
        print("connecting to redis")
        try:
            self.redis = redis.Redis(host="redis", port=6379)
        except Exception as e:
            print("failed to load redis", e)
            raise e

    def get_object(self, key: UID) -> Optional[StorableObject]:
        try:
            return self.__getitem__(key)
        except KeyError as e:  # noqa: F841
            return None

    def get_objects_of_type(self, obj_type: type) -> Iterable[StorableObject]:
        # raise NotImplementedError("get_objects_of_type")
        # return [obj for obj in self.values() if isinstance(obj.data, obj_type)]
        return self.values()

    def __sizeof__(self) -> int:
        return self.values().__sizeof__()

    def __str__(self) -> str:
        return f"{type(self)}"

    def __len__(self) -> int:
        return self.client.dbsize()

    def keys(self) -> KeysView[UID]:
        key_bytes = self.redis.keys()
        key_ids = [UID.from_string(str(key.decode("utf-8"))) for key in key_bytes]
        return key_ids

    def values(self) -> List[StorableObject]:
        key_bytes = self.redis.keys()
        # this is bad we need to decouple getting the data from the search
        all_values = []
        for key in key_bytes:
            all_values.append(self.__getitem__(key))

        # return self.redis.mget(self.keys())
        return all_values

    def __contains__(self, key: UID) -> bool:
        return self.redis.contains(str(key.value))

    def __getitem__(self, key: Union[UID, str, bytes]) -> StorableObject:
        local_session = sessionmaker(bind=self.db)()

        if isinstance(key, UID):
            key_str = str(key.value)
            key_uid = key
        elif isinstance(key, bytes):
            key_str = str(key.decode("utf-8"))
            key_uid = UID.from_string(key_str)
        else:
            key_str = key
            key_uid = UID.from_string(key_str)

        bin = self.redis.get(key_str)
        obj = syft.deserialize(bin, from_bytes=True)
        obj_metadata = (
            local_session.query(ObjectMetadata).filter_by(obj=key_str).first()
        )

        if obj is None or obj_metadata is None:
            raise KeyError(f"Object not found! for UID: {key_uid}")

        obj = StorableObject(
            id=key_uid,
            data=obj,
            description=obj_metadata.description,
            tags=obj_metadata.tags,
            read_permissions=dict(
                syft.deserialize(
                    bytes.fromhex(obj_metadata.read_permissions), from_bytes=True
                )
            ),
            search_permissions=dict(
                syft.deserialize(
                    bytes.fromhex(obj_metadata.search_permissions), from_bytes=True
                )
            ),
            # name=obj_metadata.name,
        )
        local_session.close()
        return obj

    def is_dataset(self, key: UID) -> bool:
        local_session = sessionmaker(bind=self.db)()
        is_dataset_obj = (
            local_session.query(BinObjDataset).filter_by(obj=str(key.value)).exists()
        )
        is_dataset_obj = local_session.query(is_dataset_obj).scalar()
        local_session.close()
        return is_dataset_obj

    def _get_obj_dataset_relation(self, key: UID) -> Optional[BinObjDataset]:
        local_session = sessionmaker(bind=self.db)()
        obj_dataset_relation = (
            local_session.query(BinObjDataset).filter_by(obj=str(key.value)).first()
        )
        local_session.close()
        return obj_dataset_relation

    def __setitem__(self, key: UID, value: StorableObject) -> None:

        bin = syft.serialize(value.data, to_bytes=True)
        self.redis.set(str(key.value), bin)

        metadata_obj = ObjectMetadata(
            obj=str(key.value),
            tags=value.tags,
            description=value.description,
            read_permissions=cast(
                bytes,
                syft.serialize(
                    syft.lib.python.Dict(value.read_permissions), to_bytes=True
                ),
            ).hex(),
            search_permissions=cast(
                bytes,
                syft.serialize(
                    syft.lib.python.Dict(value.search_permissions), to_bytes=True
                ),
            ).hex(),
            # name=metadata_dict["name"],
        )

        obj_dataset_relation = self._get_obj_dataset_relation(key)
        if obj_dataset_relation:
            # Create a object dataset relationship for the new object
            obj_dataset_relation = BinObjDataset(
                id=obj_dataset_relation.id,
                name=obj_dataset_relation.name,
                obj=key,
                dataset=obj_dataset_relation.dataset,
                dtype=obj_dataset_relation.dtype,
                shape=obj_dataset_relation.shape,
            )

        local_session = sessionmaker(bind=self.db)()
        local_session.add(metadata_obj)
        local_session.add(obj_dataset_relation) if obj_dataset_relation else None
        local_session.commit()
        local_session.close()

    def delete(self, key: UID) -> None:
        try:
            self.redis.delete(str(key.value))
            local_session = sessionmaker(bind=self.db)()
            metadata_to_delete = (
                local_session.query(ObjectMetadata)
                .filter_by(obj=str(key.value))
                .first()
            )
            local_session.delete(metadata_to_delete)
            local_session.commit()
            local_session.close()
        except Exception as e:
            print(f"{type(self)} Exception in __delitem__ error {key}. {e}")

    def clear(self) -> None:
        self.redis.flushdb()
        local_session = sessionmaker(bind=self.db)()
        local_session.query(ObjectMetadata).delete()
        local_session.commit()
        local_session.close()

    def __repr__(self) -> str:
        return f"{type(self)}"
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import mock
from io import BytesIO
import datetime

from django.core.management import call_command
from django.test import TestCase

class RunProtgresDatabaseBackupTest(TestCase):

    # Test the command without an argument
    @mock.patch('seqr.management.commands.run_postgres_database_backup.os')
    @mock.patch('seqr.management.commands.run_postgres_database_backup.datetime')
    def test_command_no_argument(self, mock_datetime, mock_os):
        out = BytesIO()

        mock_os.path.isdir.return_value = False
        mock_os.environ.get.side_effect = ['db_back_bucket', 'localhost', 'unknown']
        mock_datetime.datetime.now.return_value = datetime.datetime(2020, 4, 27, 20, 16, 01)
        call_command('run_postgres_database_backup', stdout = out)

        mock_os.mkdir.assert_called_with('/postgres_backups')
        calls = [
            mock.call('DATABASE_BACKUP_BUCKET', 'unknown'),
            mock.call('POSTGRES_SERVICE_HOSTNAME', 'unknown'),
            mock.call('DEPLOYMENT_TYPE', 'unknown'),
        ]
        mock_os.environ.get.assert_has_calls(calls)
        self.assertEqual(
            '=====================================\n' +
            '======== 2020-04-27__20-16-01 ======= \n' +
            '=====================================\n' +
            'Creating directory: /postgres_backups\n' +
            '/usr/bin/pg_dump -U postgres --host localhost seqrdb | gzip -c - > /postgres_backups/seqrdb_unknown_backup_2020-04-27__20-16-01.txt.gz\n' +
            '/usr/local/bin/gsutil mv /postgres_backups/seqrdb_unknown_backup_2020-04-27__20-16-01.txt.gz gs://db_back_bucket/postgres/seqrdb_unknown_backup_2020-04-27__20-16-01.txt.gz\n' +
            '/usr/bin/pg_dump -U postgres --host localhost xwiki | gzip -c - > /postgres_backups/xwiki_unknown_backup_2020-04-27__20-16-01.txt.gz\n' +
            '/usr/local/bin/gsutil mv /postgres_backups/xwiki_unknown_backup_2020-04-27__20-16-01.txt.gz gs://db_back_bucket/postgres/xwiki_unknown_backup_2020-04-27__20-16-01.txt.gz\n',
            out.getvalue())

        mock_os.system.assert_called_with(
            '/usr/local/bin/gsutil mv /postgres_backups/xwiki_unknown_backup_2020-04-27__20-16-01.txt.gz gs://db_back_bucket/postgres/xwiki_unknown_backup_2020-04-27__20-16-01.txt.gz')

    # Test the command with different arguments
    @mock.patch('seqr.management.commands.run_postgres_database_backup.os')
    @mock.patch('seqr.management.commands.run_postgres_database_backup.datetime')
    def test_command_with_arguments(self, mock_datetime, mock_os):
        out = BytesIO()

        mock_os.path.isdir.return_value = True
        mock_datetime.datetime.now.return_value = datetime.datetime(2020, 4, 27, 20, 16, 01)
        call_command('run_postgres_database_backup', '--bucket=test_bucket',
                     '--postgres-host=test_host',
                     '--deployment-type=test_deployment',
                     stdout = out)

        mock_os.mkdir.assert_not_called()
        self.assertEqual(
            '=====================================\n' +
            '======== 2020-04-27__20-16-01 ======= \n' +
            '=====================================\n' +
            '/usr/bin/pg_dump -U postgres --host test_host seqrdb | gzip -c - > /postgres_backups/seqrdb_test_deployment_backup_2020-04-27__20-16-01.txt.gz\n' +
            '/usr/local/bin/gsutil mv /postgres_backups/seqrdb_test_deployment_backup_2020-04-27__20-16-01.txt.gz gs://test_bucket/postgres/seqrdb_test_deployment_backup_2020-04-27__20-16-01.txt.gz\n' +
            '/usr/bin/pg_dump -U postgres --host test_host xwiki | gzip -c - > /postgres_backups/xwiki_test_deployment_backup_2020-04-27__20-16-01.txt.gz\n' +
            '/usr/local/bin/gsutil mv /postgres_backups/xwiki_test_deployment_backup_2020-04-27__20-16-01.txt.gz gs://test_bucket/postgres/xwiki_test_deployment_backup_2020-04-27__20-16-01.txt.gz\n',
            out.getvalue())

        mock_os.system.assert_called_with(
            '/usr/local/bin/gsutil mv /postgres_backups/xwiki_test_deployment_backup_2020-04-27__20-16-01.txt.gz gs://test_bucket/postgres/xwiki_test_deployment_backup_2020-04-27__20-16-01.txt.gz')
<commit_msg>Moved the mock.patch to ahead of the class definition.<commit_after># -*- coding: utf-8 -*-
import mock
from io import BytesIO
import datetime

from django.core.management import call_command
from django.test import TestCase

@mock.patch('seqr.management.commands.run_postgres_database_backup.os')
@mock.patch('seqr.management.commands.run_postgres_database_backup.datetime')
class RunProtgresDatabaseBackupTest(TestCase):

    # Test the command without an argument
    def test_command_no_argument(self, mock_datetime, mock_os):
        out = BytesIO()

        mock_os.path.isdir.return_value = False
        mock_os.environ.get.side_effect = ['db_back_bucket', 'localhost', 'unknown']
        mock_datetime.datetime.now.return_value = datetime.datetime(2020, 4, 27, 20, 16, 01)
        call_command('run_postgres_database_backup', stdout = out)

        mock_os.mkdir.assert_called_with('/postgres_backups')
        calls = [
            mock.call('DATABASE_BACKUP_BUCKET', 'unknown'),
            mock.call('POSTGRES_SERVICE_HOSTNAME', 'unknown'),
            mock.call('DEPLOYMENT_TYPE', 'unknown'),
        ]
        mock_os.environ.get.assert_has_calls(calls)
        self.assertEqual(
            '=====================================\n' +
            '======== 2020-04-27__20-16-01 ======= \n' +
            '=====================================\n' +
            'Creating directory: /postgres_backups\n' +
            '/usr/bin/pg_dump -U postgres --host localhost seqrdb | gzip -c - > /postgres_backups/seqrdb_unknown_backup_2020-04-27__20-16-01.txt.gz\n' +
            '/usr/local/bin/gsutil mv /postgres_backups/seqrdb_unknown_backup_2020-04-27__20-16-01.txt.gz gs://db_back_bucket/postgres/seqrdb_unknown_backup_2020-04-27__20-16-01.txt.gz\n' +
            '/usr/bin/pg_dump -U postgres --host localhost xwiki | gzip -c - > /postgres_backups/xwiki_unknown_backup_2020-04-27__20-16-01.txt.gz\n' +
            '/usr/local/bin/gsutil mv /postgres_backups/xwiki_unknown_backup_2020-04-27__20-16-01.txt.gz gs://db_back_bucket/postgres/xwiki_unknown_backup_2020-04-27__20-16-01.txt.gz\n',
            out.getvalue())

        mock_os.system.assert_called_with(
            '/usr/local/bin/gsutil mv /postgres_backups/xwiki_unknown_backup_2020-04-27__20-16-01.txt.gz gs://db_back_bucket/postgres/xwiki_unknown_backup_2020-04-27__20-16-01.txt.gz')

    # Test the command with different arguments
    def test_command_with_arguments(self, mock_datetime, mock_os):
        out = BytesIO()

        mock_os.path.isdir.return_value = True
        mock_datetime.datetime.now.return_value = datetime.datetime(2020, 4, 27, 20, 16, 01)
        call_command('run_postgres_database_backup', '--bucket=test_bucket',
                     '--postgres-host=test_host',
                     '--deployment-type=test_deployment',
                     stdout = out)

        mock_os.mkdir.assert_not_called()
        self.assertEqual(
            '=====================================\n' +
            '======== 2020-04-27__20-16-01 ======= \n' +
            '=====================================\n' +
            '/usr/bin/pg_dump -U postgres --host test_host seqrdb | gzip -c - > /postgres_backups/seqrdb_test_deployment_backup_2020-04-27__20-16-01.txt.gz\n' +
            '/usr/local/bin/gsutil mv /postgres_backups/seqrdb_test_deployment_backup_2020-04-27__20-16-01.txt.gz gs://test_bucket/postgres/seqrdb_test_deployment_backup_2020-04-27__20-16-01.txt.gz\n' +
            '/usr/bin/pg_dump -U postgres --host test_host xwiki | gzip -c - > /postgres_backups/xwiki_test_deployment_backup_2020-04-27__20-16-01.txt.gz\n' +
            '/usr/local/bin/gsutil mv /postgres_backups/xwiki_test_deployment_backup_2020-04-27__20-16-01.txt.gz gs://test_bucket/postgres/xwiki_test_deployment_backup_2020-04-27__20-16-01.txt.gz\n',
            out.getvalue())

        mock_os.system.assert_called_with(
            '/usr/local/bin/gsutil mv /postgres_backups/xwiki_test_deployment_backup_2020-04-27__20-16-01.txt.gz gs://test_bucket/postgres/xwiki_test_deployment_backup_2020-04-27__20-16-01.txt.gz')
<|endoftext|>


<commit_before># coding=utf-8
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.

"""Helper functions to use the Lithium reducer.
"""

import gzip
import io
from pathlib import Path
import re
from shlex import quote
import shutil
import subprocess
import sys
import tempfile

from lithium.interestingness.utils import file_contains_str

from . import file_manipulation
from ..js.inspect_shell import testJsShellOrXpcshell
from ..js.js_interesting import JS_OVERALL_MISMATCH
from ..js.js_interesting import JS_VG_AMISS

runlithiumpy = [sys.executable, "-u", "-m", "lithium"]  # pylint: disable=invalid-name

# Status returns for runLithium and many_timed_runs
(HAPPY, LITH_NO_REPRO, LITH_FINISHED, LITH_RETESTED_STILL_INTERESTING, LITH_BUSTED) = range(5)

def pinpoint(itest, logPrefix, jsEngine, engineFlags, infilename,  # pylint: disable=invalid-name,missing-param-doc
             bisectRepo, build_options_str, targetTime, suspiciousLevel):
    # pylint: disable=missing-return-doc,missing-return-type-doc,missing-type-doc,too-many-arguments,too-many-locals
    """Run Lithium and autobisectjs.

    itest must be an array of the form [module, ...] where module is an interestingness module.
    The module's "interesting" function must accept [...] + [jsEngine] + engineFlags + infilename
    (If it's not prepared to accept engineFlags, engineFlags must be empty.)
    """
    lithArgs = itest + [str(jsEngine)] + engineFlags + [str(infilename)]  # pylint: disable=invalid-name

    (lithResult, lithDetails) = reduction_strat(  # pylint: disable=invalid-name
        logPrefix, infilename, lithArgs, targetTime, suspiciousLevel)

    print()
    print("Done running Lithium on the part in between DDBEGIN and DDEND. To reproduce, run:")
    print(" ".join(quote(str(x)) for x in [sys.executable, "-u", "-m", "lithium", "--strategy=check-only"] + lithArgs))
    print()

    if (bisectRepo != "none" and targetTime >= 3 * 60 * 60 and
            build_options_str is not None and testJsShellOrXpcshell(jsEngine) != "xpcshell"):
        autobisectCmd = (  # pylint: disable=invalid-name
            [sys.executable, "-u", "-m", "funfuzz.autobisectjs"] +
            ["-b", build_options_str] +
            ["-p", " ".join(engineFlags + [str(infilename)])] +
            ["-i"] + [str(x) for x in itest]
        )
        print(" ".join(quote(str(x)) for x in autobisectCmd))
        autobisect_log = (logPrefix.parent / f"{logPrefix.stem}-autobisect").with_suffix(".txt")
        with io.open(str(autobisect_log), "w", encoding="utf-8", errors="replace") as f:
            subprocess.run(autobisectCmd, check=False, stderr=subprocess.STDOUT, stdout=f)
        print(f"Done running autobisectjs. Log: {autobisect_log}")

        with io.open(str(autobisect_log), "r", encoding="utf-8", errors="replace") as f:
            lines = f.readlines()
            autobisect_log_trunc = file_manipulation.truncateMid(lines, 50, ["..."])
    else:
        autobisect_log_trunc = []

    return lithResult, lithDetails, autobisect_log_trunc

def run_lithium(lithArgs, logPrefix, targetTime):  # pylint: disable=invalid-name,missing-param-doc,missing-return-doc
    # pylint: disable=missing-return-type-doc,missing-type-doc
    """Run Lithium as a subprocess: reduce to the smallest file that has at least the same unhappiness level.

    Returns a tuple of (lithlogfn, LITH_*, details).
    """
    deletableLithTemp = None  # pylint: disable=invalid-name
    if targetTime:
        # loop is being used by bot
        deletableLithTemp = tempfile.mkdtemp(prefix="fuzzbot-lithium")  # pylint: disable=invalid-name
        lithArgs = [f"--maxruntime={targetTime}", f"--tempdir={deletableLithTemp}"] + lithArgs
    else:
        # loop is being run standalone
        lithtmp = logPrefix.parent / f"{logPrefix.stem}-lith-tmp"
        Path.mkdir(lithtmp)
        lithArgs = [f"--tempdir={lithtmp}"] + lithArgs
    lithlogfn = (logPrefix.parent / f"{logPrefix.stem}-lith-out").with_suffix(".txt")
    print(f"Preparing to run Lithium, log file {lithlogfn}")
    print(" ".join(quote(str(x)) for x in runlithiumpy + lithArgs))
    with io.open(str(lithlogfn), "w", encoding="utf-8", errors="replace") as f:
        subprocess.run(runlithiumpy + lithArgs, check=False, stderr=subprocess.STDOUT, stdout=f)
    print("Done running Lithium")
    if deletableLithTemp:
        shutil.rmtree(deletableLithTemp)
    r = readLithiumResult(lithlogfn)  # pylint: disable=invalid-name

    with open(lithlogfn, "rb") as f_in:  # Replace the old gzip subprocess call
        with gzip.open(lithlogfn.with_suffix(".txt.gz"), "wb") as f_out:
            shutil.copyfileobj(f_in, f_out)
    lithlogfn.unlink()

    return r

def readLithiumResult(lithlogfn):  # pylint: disable=invalid-name,missing-docstring,missing-return-doc
    # pylint: disable=missing-return-type-doc
    with io.open(str(lithlogfn), "r", encoding="utf-8", errors="replace") as f:
        for line in f:
            if line.startswith("Lithium result"):
                print(line.rstrip())
            if line.startswith("Lithium result: interesting"):  # pylint: disable=no-else-return
                return LITH_RETESTED_STILL_INTERESTING, None
            elif line.startswith("Lithium result: succeeded, reduced to: "):
                # pylint: disable=invalid-name
                reducedTo = line[len("Lithium result: succeeded, reduced to: "):].rstrip()  # e.g. "4 lines"
                return LITH_FINISHED, reducedTo
            elif (line.startswith("Lithium result: not interesting") or
                  line.startswith("Lithium result: the original testcase is not")):
                return LITH_NO_REPRO, None
        return LITH_BUSTED, None

def reduction_strat(logPrefix, infilename, lithArgs, targetTime, lev):  # pylint: disable=invalid-name
    # pylint: disable=missing-param-doc,missing-return-doc,missing-return-type-doc,missing-type-doc,too-complex
    # pylint: disable=too-many-branches,too-many-locals,too-many-statements
    """Reduce jsfunfuzz output files using Lithium by using various strategies."""

    # This is an array because Python does not like assigning to upvars.
    reductionCount = [0]  # pylint: disable=invalid-name

    def lith_reduce(strategy):
        """Lithium reduction commands accepting various strategies.

        Args:
            strategy (str): Intended strategy to use

        Returns:
            (tuple): The finished Lithium run result and details
        """
        reductionCount[0] += 1
        # Remove empty elements
        full_lith_args = [x for x in (strategy + lithArgs) if x]
        print(" ".join(quote(str(x)) for x in [sys.executable, "-u", "-m", "lithium"] + full_lith_args))

        desc = "-chars" if strategy == "--char" else "-lines"
        (lith_result, lith_details) = run_lithium(
            full_lith_args, (logPrefix.parent / f"{logPrefix.stem}-{reductionCount[0]}{desc}"), targetTime)

        return lith_result, lith_details

    print()
    print("Running the first line reduction...")
    print()
    # Step 1: Run the first instance of line reduction.
    lith_result, lith_details = lith_reduce([])

    origNumOfLines = None  # pylint: disable=invalid-name
    if lith_details is not None:  # lith_details can be None if testcase no longer becomes interesting
        origNumOfLines = int(lith_details.split()[0])  # pylint: disable=invalid-name

    hasTryItOut = False  # pylint: disable=invalid-name
    hasTryItOutRegex = re.compile(r'count=[0-9]+; tryItOut\("')  # pylint: disable=invalid-name

    with io.open(str(infilename), "r", encoding="utf-8", errors="replace") as f:
        for line in file_manipulation.linesWith(f, '; tryItOut("'):
            # Checks if testcase came from jsfunfuzz or compare_jit.
            # Do not use .match here, it only matches from the start of the line:
            # https://docs.python.org/2/library/re.html#search-vs-match
            hasTryItOut = hasTryItOutRegex.search(line)  # pylint: disable=invalid-name
            if hasTryItOut:  # Stop searching after finding the first tryItOut line.
                break

    # Step 2: Run 1 instance of 1-line reduction after moving tryItOut and count=X around.
    if lith_result == LITH_FINISHED and origNumOfLines <= 50 and hasTryItOut and lev >= JS_VG_AMISS:

        tryItOutAndCountRegex = re.compile(r'"\);\ncount=([0-9]+); tryItOut\("',  # pylint: disable=invalid-name
                                           re.MULTILINE)
        with io.open(str(infilename), "r", encoding="utf-8", errors="replace") as f:
            infileContents = f.read()  # pylint: disable=invalid-name
            infileContents = re.sub(tryItOutAndCountRegex,  # pylint: disable=invalid-name
                                    ';\\\n"); count=\\1; tryItOut("\\\n',
                                    infileContents)
        with io.open(str(infilename), "w", encoding="utf-8", errors="replace") as f:
            f.write(infileContents)

        print()
        print("Running 1 instance of 1-line reduction after moving tryItOut and count=X...")
        print()
        # --chunksize=1: Reduce only individual lines, for only 1 round.
        lith_result, lith_details = lith_reduce(["--chunksize=1"])

    # Step 3: Run 1 instance of 2-line reduction after moving count=X to its own line and add a
    # 1-line offset.
    if lith_result == LITH_FINISHED and origNumOfLines <= 50 and hasTryItOut and lev >= JS_VG_AMISS:
        intendedLines = []  # pylint: disable=invalid-name
        with io.open(str(infilename), "r", encoding="utf-8", errors="replace") as f:
            for line in f:  # The testcase is likely to already be partially reduced.
                if "dumpln(cookie" not in line:  # jsfunfuzz-specific line ignore
                    # This should be simpler than re.compile.
                    intendedLines.append(line.replace("; count=", ";\ncount=")
                                         .replace('; tryItOut("', ';\ntryItOut("')
                                         # The 1-line offset is added here.
                                         .replace("SPLICE DDBEGIN", "SPLICE DDBEGIN\n"))

        with io.open(str(infilename), "w", encoding="utf-8", errors="replace") as f:
            f.writelines(intendedLines)
        print()
        print("Running 1 instance of 2-line reduction after moving count=X to its own line...")
        print()
        lith_result, lith_details = lith_reduce(["--chunksize=2"])

    # Step 4: Run 1 instance of 2-line reduction again, e.g. to remove pairs of STRICT_MODE lines.
    if lith_result == LITH_FINISHED and origNumOfLines <= 50 and hasTryItOut and lev >= JS_VG_AMISS:
        print()
        print("Running 1 instance of 2-line reduction again...")
        print()
        lith_result, lith_details = lith_reduce(["--chunksize=2"])

    isLevOverallMismatchAsmJsAvailable = (lev == JS_OVERALL_MISMATCH and  # pylint: disable=invalid-name
                                          file_contains_str(str(infilename), b"isAsmJSCompilationAvailable"))
    # Step 5 (not always run): Run character reduction within interesting lines.
    if lith_result == LITH_FINISHED and origNumOfLines <= 50 and targetTime is None and \
            lev >= JS_OVERALL_MISMATCH and not isLevOverallMismatchAsmJsAvailable:
        print()
        print("Running character reduction...")
        print()
        lith_result, lith_details = lith_reduce(["--char"])

    # Step 6: Run line reduction after activating SECOND DDBEGIN with a 1-line offset.
    if lith_result == LITH_FINISHED and origNumOfLines <= 50 and hasTryItOut and lev >= JS_VG_AMISS:
        infileContents = []  # pylint: disable=invalid-name
        with io.open(str(infilename), "r", encoding="utf-8", errors="replace") as f:
            for line in f:
                if "NIGEBDD" in line:
                    infileContents.append(line.replace("NIGEBDD", "DDBEGIN"))
                    infileContents.append("\n")  # The 1-line offset is added here.
                    continue
                infileContents.append(line)
        with io.open(str(infilename), "w", encoding="utf-8", errors="replace") as f:
            f.writelines(infileContents)

        print()
        print("Running line reduction with a 1-line offset...")
        print()
        lith_result, lith_details = lith_reduce([])

    # Step 7: Run line reduction for a final time.
    if lith_result == LITH_FINISHED and origNumOfLines <= 50 and hasTryItOut and lev >= JS_VG_AMISS:
        print()
        print("Running the final line reduction...")
        print()
        lith_result, lith_details = lith_reduce([])

    return lith_result, lith_details
<commit_msg>Fix lithium_helpers.pinpoint when running with --repo=none<commit_after># coding=utf-8
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.

"""Helper functions to use the Lithium reducer.
"""

import gzip
import io
from pathlib import Path
import re
from shlex import quote
import shutil
import subprocess
import sys
import tempfile

from lithium.interestingness.utils import file_contains_str

from . import file_manipulation
from ..js.inspect_shell import testJsShellOrXpcshell
from ..js.js_interesting import JS_OVERALL_MISMATCH
from ..js.js_interesting import JS_VG_AMISS

runlithiumpy = [sys.executable, "-u", "-m", "lithium"]  # pylint: disable=invalid-name

# Status returns for runLithium and many_timed_runs
(HAPPY, LITH_NO_REPRO, LITH_FINISHED, LITH_RETESTED_STILL_INTERESTING, LITH_BUSTED) = range(5)

def pinpoint(itest, logPrefix, jsEngine, engineFlags, infilename,  # pylint: disable=invalid-name,missing-param-doc
             bisectRepo, build_options_str, targetTime, suspiciousLevel):
    # pylint: disable=missing-return-doc,missing-return-type-doc,missing-type-doc,too-many-arguments,too-many-locals
    """Run Lithium and autobisectjs.

    itest must be an array of the form [module, ...] where module is an interestingness module.
    The module's "interesting" function must accept [...] + [jsEngine] + engineFlags + infilename
    (If it's not prepared to accept engineFlags, engineFlags must be empty.)
    """
    lithArgs = itest + [str(jsEngine)] + engineFlags + [str(infilename)]  # pylint: disable=invalid-name

    (lithResult, lithDetails) = reduction_strat(  # pylint: disable=invalid-name
        logPrefix, infilename, lithArgs, targetTime, suspiciousLevel)

    print()
    print("Done running Lithium on the part in between DDBEGIN and DDEND. To reproduce, run:")
    print(" ".join(quote(str(x)) for x in [sys.executable, "-u", "-m", "lithium", "--strategy=check-only"] + lithArgs))
    print()

    if (bisectRepo != Path("none") and targetTime >= 3 * 60 * 60 and
            build_options_str is not None and testJsShellOrXpcshell(jsEngine) != "xpcshell"):
        autobisectCmd = (  # pylint: disable=invalid-name
            [sys.executable, "-u", "-m", "funfuzz.autobisectjs"] +
            ["-b", build_options_str] +
            ["-p", " ".join(engineFlags + [str(infilename)])] +
            ["-i"] + [str(x) for x in itest]
        )
        print(" ".join(quote(str(x)) for x in autobisectCmd))
        autobisect_log = (logPrefix.parent / f"{logPrefix.stem}-autobisect").with_suffix(".txt")
        with io.open(str(autobisect_log), "w", encoding="utf-8", errors="replace") as f:
            subprocess.run(autobisectCmd, check=False, stderr=subprocess.STDOUT, stdout=f)
        print(f"Done running autobisectjs. Log: {autobisect_log}")

        with io.open(str(autobisect_log), "r", encoding="utf-8", errors="replace") as f:
            lines = f.readlines()
            autobisect_log_trunc = file_manipulation.truncateMid(lines, 50, ["..."])
    else:
        autobisect_log_trunc = []

    return lithResult, lithDetails, autobisect_log_trunc

def run_lithium(lithArgs, logPrefix, targetTime):  # pylint: disable=invalid-name,missing-param-doc,missing-return-doc
    # pylint: disable=missing-return-type-doc,missing-type-doc
    """Run Lithium as a subprocess: reduce to the smallest file that has at least the same unhappiness level.

    Returns a tuple of (lithlogfn, LITH_*, details).
    """
    deletableLithTemp = None  # pylint: disable=invalid-name
    if targetTime:
        # loop is being used by bot
        deletableLithTemp = tempfile.mkdtemp(prefix="fuzzbot-lithium")  # pylint: disable=invalid-name
        lithArgs = [f"--maxruntime={targetTime}", f"--tempdir={deletableLithTemp}"] + lithArgs
    else:
        # loop is being run standalone
        lithtmp = logPrefix.parent / f"{logPrefix.stem}-lith-tmp"
        Path.mkdir(lithtmp)
        lithArgs = [f"--tempdir={lithtmp}"] + lithArgs
    lithlogfn = (logPrefix.parent / f"{logPrefix.stem}-lith-out").with_suffix(".txt")
    print(f"Preparing to run Lithium, log file {lithlogfn}")
    print(" ".join(quote(str(x)) for x in runlithiumpy + lithArgs))
    with io.open(str(lithlogfn), "w", encoding="utf-8", errors="replace") as f:
        subprocess.run(runlithiumpy + lithArgs, check=False, stderr=subprocess.STDOUT, stdout=f)
    print("Done running Lithium")
    if deletableLithTemp:
        shutil.rmtree(deletableLithTemp)
    r = readLithiumResult(lithlogfn)  # pylint: disable=invalid-name

    with open(lithlogfn, "rb") as f_in:  # Replace the old gzip subprocess call
        with gzip.open(lithlogfn.with_suffix(".txt.gz"), "wb") as f_out:
            shutil.copyfileobj(f_in, f_out)
    lithlogfn.unlink()

    return r

def readLithiumResult(lithlogfn):  # pylint: disable=invalid-name,missing-docstring,missing-return-doc
    # pylint: disable=missing-return-type-doc
    with io.open(str(lithlogfn), "r", encoding="utf-8", errors="replace") as f:
        for line in f:
            if line.startswith("Lithium result"):
                print(line.rstrip())
            if line.startswith("Lithium result: interesting"):  # pylint: disable=no-else-return
                return LITH_RETESTED_STILL_INTERESTING, None
            elif line.startswith("Lithium result: succeeded, reduced to: "):
                # pylint: disable=invalid-name
                reducedTo = line[len("Lithium result: succeeded, reduced to: "):].rstrip()  # e.g. "4 lines"
                return LITH_FINISHED, reducedTo
            elif (line.startswith("Lithium result: not interesting") or
                  line.startswith("Lithium result: the original testcase is not")):
                return LITH_NO_REPRO, None
        return LITH_BUSTED, None

def reduction_strat(logPrefix, infilename, lithArgs, targetTime, lev):  # pylint: disable=invalid-name
    # pylint: disable=missing-param-doc,missing-return-doc,missing-return-type-doc,missing-type-doc,too-complex
    # pylint: disable=too-many-branches,too-many-locals,too-many-statements
    """Reduce jsfunfuzz output files using Lithium by using various strategies."""

    # This is an array because Python does not like assigning to upvars.
    reductionCount = [0]  # pylint: disable=invalid-name

    def lith_reduce(strategy):
        """Lithium reduction commands accepting various strategies.

        Args:
            strategy (str): Intended strategy to use

        Returns:
            (tuple): The finished Lithium run result and details
        """
        reductionCount[0] += 1
        # Remove empty elements
        full_lith_args = [x for x in (strategy + lithArgs) if x]
        print(" ".join(quote(str(x)) for x in [sys.executable, "-u", "-m", "lithium"] + full_lith_args))

        desc = "-chars" if strategy == "--char" else "-lines"
        (lith_result, lith_details) = run_lithium(
            full_lith_args, (logPrefix.parent / f"{logPrefix.stem}-{reductionCount[0]}{desc}"), targetTime)

        return lith_result, lith_details

    print()
    print("Running the first line reduction...")
    print()
    # Step 1: Run the first instance of line reduction.
    lith_result, lith_details = lith_reduce([])

    origNumOfLines = None  # pylint: disable=invalid-name
    if lith_details is not None:  # lith_details can be None if testcase no longer becomes interesting
        origNumOfLines = int(lith_details.split()[0])  # pylint: disable=invalid-name

    hasTryItOut = False  # pylint: disable=invalid-name
    hasTryItOutRegex = re.compile(r'count=[0-9]+; tryItOut\("')  # pylint: disable=invalid-name

    with io.open(str(infilename), "r", encoding="utf-8", errors="replace") as f:
        for line in file_manipulation.linesWith(f, '; tryItOut("'):
            # Checks if testcase came from jsfunfuzz or compare_jit.
            # Do not use .match here, it only matches from the start of the line:
            # https://docs.python.org/2/library/re.html#search-vs-match
            hasTryItOut = hasTryItOutRegex.search(line)  # pylint: disable=invalid-name
            if hasTryItOut:  # Stop searching after finding the first tryItOut line.
                break

    # Step 2: Run 1 instance of 1-line reduction after moving tryItOut and count=X around.
    if lith_result == LITH_FINISHED and origNumOfLines <= 50 and hasTryItOut and lev >= JS_VG_AMISS:

        tryItOutAndCountRegex = re.compile(r'"\);\ncount=([0-9]+); tryItOut\("',  # pylint: disable=invalid-name
                                           re.MULTILINE)
        with io.open(str(infilename), "r", encoding="utf-8", errors="replace") as f:
            infileContents = f.read()  # pylint: disable=invalid-name
            infileContents = re.sub(tryItOutAndCountRegex,  # pylint: disable=invalid-name
                                    ';\\\n"); count=\\1; tryItOut("\\\n',
                                    infileContents)
        with io.open(str(infilename), "w", encoding="utf-8", errors="replace") as f:
            f.write(infileContents)

        print()
        print("Running 1 instance of 1-line reduction after moving tryItOut and count=X...")
        print()
        # --chunksize=1: Reduce only individual lines, for only 1 round.
        lith_result, lith_details = lith_reduce(["--chunksize=1"])

    # Step 3: Run 1 instance of 2-line reduction after moving count=X to its own line and add a
    # 1-line offset.
    if lith_result == LITH_FINISHED and origNumOfLines <= 50 and hasTryItOut and lev >= JS_VG_AMISS:
        intendedLines = []  # pylint: disable=invalid-name
        with io.open(str(infilename), "r", encoding="utf-8", errors="replace") as f:
            for line in f:  # The testcase is likely to already be partially reduced.
                if "dumpln(cookie" not in line:  # jsfunfuzz-specific line ignore
                    # This should be simpler than re.compile.
                    intendedLines.append(line.replace("; count=", ";\ncount=")
                                         .replace('; tryItOut("', ';\ntryItOut("')
                                         # The 1-line offset is added here.
                                         .replace("SPLICE DDBEGIN", "SPLICE DDBEGIN\n"))

        with io.open(str(infilename), "w", encoding="utf-8", errors="replace") as f:
            f.writelines(intendedLines)
        print()
        print("Running 1 instance of 2-line reduction after moving count=X to its own line...")
        print()
        lith_result, lith_details = lith_reduce(["--chunksize=2"])

    # Step 4: Run 1 instance of 2-line reduction again, e.g. to remove pairs of STRICT_MODE lines.
    if lith_result == LITH_FINISHED and origNumOfLines <= 50 and hasTryItOut and lev >= JS_VG_AMISS:
        print()
        print("Running 1 instance of 2-line reduction again...")
        print()
        lith_result, lith_details = lith_reduce(["--chunksize=2"])

    isLevOverallMismatchAsmJsAvailable = (lev == JS_OVERALL_MISMATCH and  # pylint: disable=invalid-name
                                          file_contains_str(str(infilename), b"isAsmJSCompilationAvailable"))
    # Step 5 (not always run): Run character reduction within interesting lines.
    if lith_result == LITH_FINISHED and origNumOfLines <= 50 and targetTime is None and \
            lev >= JS_OVERALL_MISMATCH and not isLevOverallMismatchAsmJsAvailable:
        print()
        print("Running character reduction...")
        print()
        lith_result, lith_details = lith_reduce(["--char"])

    # Step 6: Run line reduction after activating SECOND DDBEGIN with a 1-line offset.
    if lith_result == LITH_FINISHED and origNumOfLines <= 50 and hasTryItOut and lev >= JS_VG_AMISS:
        infileContents = []  # pylint: disable=invalid-name
        with io.open(str(infilename), "r", encoding="utf-8", errors="replace") as f:
            for line in f:
                if "NIGEBDD" in line:
                    infileContents.append(line.replace("NIGEBDD", "DDBEGIN"))
                    infileContents.append("\n")  # The 1-line offset is added here.
                    continue
                infileContents.append(line)
        with io.open(str(infilename), "w", encoding="utf-8", errors="replace") as f:
            f.writelines(infileContents)

        print()
        print("Running line reduction with a 1-line offset...")
        print()
        lith_result, lith_details = lith_reduce([])

    # Step 7: Run line reduction for a final time.
    if lith_result == LITH_FINISHED and origNumOfLines <= 50 and hasTryItOut and lev >= JS_VG_AMISS:
        print()
        print("Running the final line reduction...")
        print()
        lith_result, lith_details = lith_reduce([])

    return lith_result, lith_details
<|endoftext|>


<commit_before>import os
import math
from functools import partial

import numpy as np
from keras import backend as K
from keras.layers import (Conv3D, MaxPooling3D, Activation, UpSampling3D, merge, Input, Reshape)
from keras.models import Model, load_model
from keras.optimizers import SGD
from keras.callbacks import ModelCheckpoint, CSVLogger, Callback, LearningRateScheduler

from DataGenerator import get_training_and_testing_generators

pool_size = (2, 2, 2)
image_shape = (144, 240, 240)
n_channels = 3
input_shape = tuple([n_channels] + list(image_shape))
n_labels = 1  # not including background
batch_size = 1
z_crop = 155 - image_shape[0]
n_epochs = 50
data_dir = "/home/neuro-user/PycharmProjects/BRATS/sample_data"
truth_channel = 3
background_channel = 4
decay_learning_rate_every_x_epochs = 1
initial_learning_rate = 0.1
learning_rate_drop = 0.5
validation_split = 0.8

# learning rate schedule
def step_decay(epoch, initial_lrate=initial_learning_rate, drop=learning_rate_drop,
               epochs_drop=decay_learning_rate_every_x_epochs):
    return initial_lrate * math.pow(drop, math.floor((1+epoch)/float(epochs_drop)))

K.set_image_dim_ordering('th')
smooth = 1.

def dice_coef(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)

def dice_coef_loss(y_true, y_pred):
    return -dice_coef(y_true, y_pred)

def unet_model():
    inputs = Input(input_shape)
    conv1 = Conv3D(32, 3, 3, 3, activation='relu', border_mode='same')(inputs)
    conv1 = Conv3D(32, 3, 3, 3, activation='relu', border_mode='same')(conv1)
    pool1 = MaxPooling3D(pool_size=pool_size)(conv1)

    conv2 = Conv3D(64, 3, 3, 3, activation='relu', border_mode='same')(pool1)
    conv2 = Conv3D(64, 3, 3, 3, activation='relu', border_mode='same')(conv2)
    pool2 = MaxPooling3D(pool_size=pool_size)(conv2)

    conv3 = Conv3D(128, 3, 3, 3, activation='relu', border_mode='same')(pool2)
    conv3 = Conv3D(128, 3, 3, 3, activation='relu', border_mode='same')(conv3)
    pool3 = MaxPooling3D(pool_size=pool_size)(conv3)

    conv4 = Conv3D(256, 3, 3, 3, activation='relu', border_mode='same')(pool3)
    conv4 = Conv3D(256, 3, 3, 3, activation='relu', border_mode='same')(conv4)
    pool4 = MaxPooling3D(pool_size=pool_size)(conv4)

    conv5 = Conv3D(512, 3, 3, 3, activation='relu', border_mode='same')(pool4)
    conv5 = Conv3D(512, 3, 3, 3, activation='relu', border_mode='same')(conv5)

    up6 = merge([UpSampling3D(size=pool_size)(conv5), conv4], mode='concat', concat_axis=1)
    conv6 = Conv3D(256, 3, 3, 3, activation='relu', border_mode='same')(up6)
    conv6 = Conv3D(256, 3, 3, 3, activation='relu', border_mode='same')(conv6)

    up7 = merge([UpSampling3D(size=pool_size)(conv6), conv3], mode='concat', concat_axis=1)
    conv7 = Conv3D(128, 3, 3, 3, activation='relu', border_mode='same')(up7)
    conv7 = Conv3D(128, 3, 3, 3, activation='relu', border_mode='same')(conv7)

    up8 = merge([UpSampling3D(size=pool_size)(conv7), conv2], mode='concat', concat_axis=1)
    conv8 = Conv3D(64, 3, 3, 3, activation='relu', border_mode='same')(up8)
    conv8 = Conv3D(64, 3, 3, 3, activation='relu', border_mode='same')(conv8)

    up9 = merge([UpSampling3D(size=pool_size)(conv8), conv1], mode='concat', concat_axis=1)
    conv9 = Conv3D(32, 3, 3, 3, activation='relu', border_mode='same')(up9)
    conv9 = Conv3D(32, 3, 3, 3, activation='relu', border_mode='same')(conv9)

    conv10 = Conv3D(n_labels, 1, 1, 1)(conv9)
    act = Activation('sigmoid')(conv10)
    model = Model(input=inputs, output=act)

    model.compile(optimizer=SGD(lr=initial_learning_rate, decay=0.0, momentum=0.9), loss=dice_coef_loss,
                  metrics=[dice_coef])

    return model

def get_training_weights(training_generator, nb_training_samples, n_classes=2):
    counts = np.zeros((n_labels, n_classes))
    i = 0
    while i < nb_training_samples:
        _, y_train = training_generator.next()
        if n_labels > 1:
            for label in range(n_labels):
                counts[label, :] += np.bincount(np.array(y_train[:, label].ravel(), dtype=np.uint8),
                                                minlength=n_classes)
        else:
            counts += np.bincount(np.array(y_train.ravel(), dtype=np.uint8), minlength=n_classes)
        i += 1
    print("Class Counts: {0}".format(counts))
    return counts_to_weights(counts)

def counts_to_weights(array):
    weights_list = []
    for label in range(array.shape[0]):
        weights = float(array[label, :].max())/array[label, :]
        weights_list.append({0: weights[0], 1: weights[1]})
    if len(weights_list) == 1:
        return weights_list[0]
    else:
        return weights_list

def deleteme(array):
    length = len(array)
    if length > 2:
        out_list = []
        array_sum = array.sum()
        for item in array:
            background_count = array_sum - item
            out_list.append({0:1, 1:float(background_count)/item})
        return out_list
    else:
        out_dict = dict()
        weights = []
        for i, item in enumerate(weights):
            out_dict[i] = item
        return out_dict

def get_training_weights_from_data(y_train):
    weights = []
    for label in range(y_train.shape[1]):
        background, foreground = get_class_weights(y_train[:, label], n_classes=2)
        weights.append({0: background, 1: foreground})
    print("Training weights: {0}".format(weights))
    if len(weights) == 1:
        return weights[0]
    else:
        return weights

def get_class_weights(labels_array, n_classes):
    counts = np.bincount(np.array(labels_array.ravel(), dtype=np.uint8), minlength=n_classes)
    weights = np.zeros(n_classes)
    weights[counts > 0] = float(counts.max())/counts[counts > 0]
    return weights

class SaveLossHistory(Callback):
    def on_train_begin(self, logs={}):
        self.losses = []

    def on_batch_end(self, batch, logs={}):
        self.losses.append(logs.get('loss'))
        pickle_dump(self.losses, "loss_history.pkl")

def get_callbacks(model_file):
    epoch_model_file = model_file.replace(".h", "_latest.h")
    model_checkpoint = ModelCheckpoint(epoch_model_file)
    logger = CSVLogger("training.log")
    history = SaveLossHistory()
    scheduler = LearningRateScheduler(partial(step_decay, initial_lrate=initial_learning_rate, drop=learning_rate_drop,
                                              epochs_drop=decay_learning_rate_every_x_epochs))
    return [model_checkpoint, logger, history, scheduler]

def main(overwrite=False):
    model_file = os.path.abspath("3d_unet_model.h5")
    if not overwrite and os.path.exists(model_file):
        print("Loading pre-trained model")
        model = load_model(model_file, custom_objects={'dice_coef_loss': dice_coef_loss, 'dice_coef': dice_coef})
    else:
        model = unet_model()
    train_model(model, model_file)

def train_model(model, model_file):
    training_generator, testing_generator, nb_training_samples, nb_testing_samples = get_training_and_testing_generators(
        data_dir=data_dir, batch_size=batch_size, nb_channels=n_channels, truth_channel=truth_channel, z_crop=z_crop,
        background_channel=background_channel, validation_split=validation_split)

    model.fit_generator(generator=training_generator,
                        samples_per_epoch=nb_training_samples,
                        nb_epoch=n_epochs,
                        validation_data=testing_generator,
                        nb_val_samples=nb_testing_samples,
                        pickle_safe=True,
                        callbacks=get_callbacks(model_file))
    model.save(model_file)

if __name__ == "__main__":
    main(overwrite=False)
<commit_msg>FIX: imports pickle_dump<commit_after>import os
import math
from functools import partial

import numpy as np
from keras import backend as K
from keras.layers import (Conv3D, MaxPooling3D, Activation, UpSampling3D, merge, Input, Reshape)
from keras.models import Model, load_model
from keras.optimizers import SGD
from keras.callbacks import ModelCheckpoint, CSVLogger, Callback, LearningRateScheduler

from DataGenerator import get_training_and_testing_generators, pickle_dump

pool_size = (2, 2, 2)
image_shape = (144, 240, 240)
n_channels = 3
input_shape = tuple([n_channels] + list(image_shape))
n_labels = 1  # not including background
batch_size = 1
z_crop = 155 - image_shape[0]
n_epochs = 50
data_dir = "/home/neuro-user/PycharmProjects/BRATS/sample_data"
truth_channel = 3
background_channel = 4
decay_learning_rate_every_x_epochs = 1
initial_learning_rate = 0.1
learning_rate_drop = 0.5
validation_split = 0.8

# learning rate schedule
def step_decay(epoch, initial_lrate=initial_learning_rate, drop=learning_rate_drop,
               epochs_drop=decay_learning_rate_every_x_epochs):
    return initial_lrate * math.pow(drop, math.floor((1+epoch)/float(epochs_drop)))

K.set_image_dim_ordering('th')
smooth = 1.

def dice_coef(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)

def dice_coef_loss(y_true, y_pred):
    return -dice_coef(y_true, y_pred)

def unet_model():
    inputs = Input(input_shape)
    conv1 = Conv3D(32, 3, 3, 3, activation='relu', border_mode='same')(inputs)
    conv1 = Conv3D(32, 3, 3, 3, activation='relu', border_mode='same')(conv1)
    pool1 = MaxPooling3D(pool_size=pool_size)(conv1)

    conv2 = Conv3D(64, 3, 3, 3, activation='relu', border_mode='same')(pool1)
    conv2 = Conv3D(64, 3, 3, 3, activation='relu', border_mode='same')(conv2)
    pool2 = MaxPooling3D(pool_size=pool_size)(conv2)

    conv3 = Conv3D(128, 3, 3, 3, activation='relu', border_mode='same')(pool2)
    conv3 = Conv3D(128, 3, 3, 3, activation='relu', border_mode='same')(conv3)
    pool3 = MaxPooling3D(pool_size=pool_size)(conv3)

    conv4 = Conv3D(256, 3, 3, 3, activation='relu', border_mode='same')(pool3)
    conv4 = Conv3D(256, 3, 3, 3, activation='relu', border_mode='same')(conv4)
    pool4 = MaxPooling3D(pool_size=pool_size)(conv4)

    conv5 = Conv3D(512, 3, 3, 3, activation='relu', border_mode='same')(pool4)
    conv5 = Conv3D(512, 3, 3, 3, activation='relu', border_mode='same')(conv5)

    up6 = merge([UpSampling3D(size=pool_size)(conv5), conv4], mode='concat', concat_axis=1)
    conv6 = Conv3D(256, 3, 3, 3, activation='relu', border_mode='same')(up6)
    conv6 = Conv3D(256, 3, 3, 3, activation='relu', border_mode='same')(conv6)

    up7 = merge([UpSampling3D(size=pool_size)(conv6), conv3], mode='concat', concat_axis=1)
    conv7 = Conv3D(128, 3, 3, 3, activation='relu', border_mode='same')(up7)
    conv7 = Conv3D(128, 3, 3, 3, activation='relu', border_mode='same')(conv7)

    up8 = merge([UpSampling3D(size=pool_size)(conv7), conv2], mode='concat', concat_axis=1)
    conv8 = Conv3D(64, 3, 3, 3, activation='relu', border_mode='same')(up8)
    conv8 = Conv3D(64, 3, 3, 3, activation='relu', border_mode='same')(conv8)

    up9 = merge([UpSampling3D(size=pool_size)(conv8), conv1], mode='concat', concat_axis=1)
    conv9 = Conv3D(32, 3, 3, 3, activation='relu', border_mode='same')(up9)
    conv9 = Conv3D(32, 3, 3, 3, activation='relu', border_mode='same')(conv9)

    conv10 = Conv3D(n_labels, 1, 1, 1)(conv9)
    act = Activation('sigmoid')(conv10)
    model = Model(input=inputs, output=act)

    model.compile(optimizer=SGD(lr=initial_learning_rate, decay=0.0, momentum=0.9), loss=dice_coef_loss,
                  metrics=[dice_coef])

    return model

def get_training_weights(training_generator, nb_training_samples, n_classes=2):
    counts = np.zeros((n_labels, n_classes))
    i = 0
    while i < nb_training_samples:
        _, y_train = training_generator.next()
        if n_labels > 1:
            for label in range(n_labels):
                counts[label, :] += np.bincount(np.array(y_train[:, label].ravel(), dtype=np.uint8),
                                                minlength=n_classes)
        else:
            counts += np.bincount(np.array(y_train.ravel(), dtype=np.uint8), minlength=n_classes)
        i += 1
    print("Class Counts: {0}".format(counts))
    return counts_to_weights(counts)

def counts_to_weights(array):
    weights_list = []
    for label in range(array.shape[0]):
        weights = float(array[label, :].max())/array[label, :]
        weights_list.append({0: weights[0], 1: weights[1]})
    if len(weights_list) == 1:
        return weights_list[0]
    else:
        return weights_list

def deleteme(array):
    length = len(array)
    if length > 2:
        out_list = []
        array_sum = array.sum()
        for item in array:
            background_count = array_sum - item
            out_list.append({0:1, 1:float(background_count)/item})
        return out_list
    else:
        out_dict = dict()
        weights = []
        for i, item in enumerate(weights):
            out_dict[i] = item
        return out_dict

def get_training_weights_from_data(y_train):
    weights = []
    for label in range(y_train.shape[1]):
        background, foreground = get_class_weights(y_train[:, label], n_classes=2)
        weights.append({0: background, 1: foreground})
    print("Training weights: {0}".format(weights))
    if len(weights) == 1:
        return weights[0]
    else:
        return weights

def get_class_weights(labels_array, n_classes):
    counts = np.bincount(np.array(labels_array.ravel(), dtype=np.uint8), minlength=n_classes)
    weights = np.zeros(n_classes)
    weights[counts > 0] = float(counts.max())/counts[counts > 0]
    return weights

class SaveLossHistory(Callback):
    def on_train_begin(self, logs={}):
        self.losses = []

    def on_batch_end(self, batch, logs={}):
        self.losses.append(logs.get('loss'))
        pickle_dump(self.losses, "loss_history.pkl")

def get_callbacks(model_file):
    epoch_model_file = model_file.replace(".h", "_latest.h")
    model_checkpoint = ModelCheckpoint(epoch_model_file)
    logger = CSVLogger("training.log")
    history = SaveLossHistory()
    scheduler = LearningRateScheduler(partial(step_decay, initial_lrate=initial_learning_rate, drop=learning_rate_drop,
                                              epochs_drop=decay_learning_rate_every_x_epochs))
    return [model_checkpoint, logger, history, scheduler]

def main(overwrite=False):
    model_file = os.path.abspath("3d_unet_model.h5")
    if not overwrite and os.path.exists(model_file):
        print("Loading pre-trained model")
        model = load_model(model_file, custom_objects={'dice_coef_loss': dice_coef_loss, 'dice_coef': dice_coef})
    else:
        model = unet_model()
    train_model(model, model_file)

def train_model(model, model_file):
    training_generator, testing_generator, nb_training_samples, nb_testing_samples = get_training_and_testing_generators(
        data_dir=data_dir, batch_size=batch_size, nb_channels=n_channels, truth_channel=truth_channel, z_crop=z_crop,
        background_channel=background_channel, validation_split=validation_split)

    model.fit_generator(generator=training_generator,
                        samples_per_epoch=nb_training_samples,
                        nb_epoch=n_epochs,
                        validation_data=testing_generator,
                        nb_val_samples=nb_testing_samples,
                        pickle_safe=True,
                        callbacks=get_callbacks(model_file))
    model.save(model_file)

if __name__ == "__main__":
    main(overwrite=False)
<|endoftext|>


<commit_before>'''
Written by Thilina Rajapakse
https://github.com/ThilinaRajapakse/BERT_binary_text_classification/blob/master/converter.py
'''

import torch, sys
import pickle
from math import ceil

from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler,
                              TensorDataset)
from torch.nn import CrossEntropyLoss, MSELoss

from tqdm import tqdm, trange
import os
from pytorch_pretrained_bert import BertTokenizer, BertModel, BertForMaskedLM, BertForSequenceClassification
from pytorch_pretrained_bert.optimization import BertAdam, WarmupLinearSchedule

from multiprocessing import Pool, cpu_count
from tools import *
import convert_examples_to_features

# OPTIONAL: if you want to have more information on what's happening, activate the logger as follows
import logging
logging.basicConfig(level=logging.INFO)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Bert pre-trained model selected in the list: bert-base-uncased,
# bert-large-uncased, bert-base-cased, bert-large-cased, bert-base-multilingual-uncased,
# bert-base-multilingual-cased, bert-base-chinese.
BERT_MODEL = 'bert-base-uncased'

# The name of the task to train.
TASK_NAME = sys.argv[1]

# The input data dir. Should contain the .tsv files (or other data files) for the task.
DATA_DIR = f'data/{TASK_NAME}/'

CACHE_DIR = 'cache/'

# The maximum total input sequence length after WordPiece tokenization.
# Sequences longer than this will be truncated, and sequences shorter than this will be padded.
MAX_SEQ_LENGTH = 128

TRAIN_BATCH_SIZE = 24
EVAL_BATCH_SIZE = 8
LEARNING_RATE = 2e-5
NUM_TRAIN_EPOCHS = 1
RANDOM_SEED = 42
GRADIENT_ACCUMULATION_STEPS = 1
WARMUP_PROPORTION = 0.1

if sys.argv[2] == 'reg':
    OUTPUT_MODE = 'regression'
    # The output directory where the model predictions and checkpoints will be written.
    OUTPUT_DIR = 'outputs/' + TASK_NAME + '_reg/'
else:
    OUTPUT_MODE = 'classification'
    # The output directory where the model predictions and checkpoints will be written.
    OUTPUT_DIR = 'outputs/' + TASK_NAME + '/'

CONFIG_NAME = "config.json"
WEIGHTS_NAME = "pytorch_model.bin"

num_labels = 2

processor = BinaryClassificationProcessor()
train_examples = processor.get_train_examples(DATA_DIR)
train_examples_len = len(train_examples)
print("Train examples count: ", train_examples_len)

with open(OUTPUT_DIR + "train_features.pkl", "rb") as f:
    train_features = pickle.load(f)

train_examples_len = len(train_features)
print("Train features count: ", train_examples_len)

num_train_optimization_steps = ceil(
    train_examples_len / TRAIN_BATCH_SIZE / GRADIENT_ACCUMULATION_STEPS) * NUM_TRAIN_EPOCHS

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = BertForSequenceClassification.from_pretrained('bert-base-uncased', cache_dir='cache', num_labels = num_labels)

model.to(device)

param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
    {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
    {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
    ]

optimizer = BertAdam(optimizer_grouped_parameters,
                     lr=LEARNING_RATE,
                     warmup=WARMUP_PROPORTION,
                     t_total=num_train_optimization_steps)

global_step = 0
nb_tr_steps = 0
tr_loss = 0

logger.info("***** Running training *****")
logger.info("  Num examples = %d", train_examples_len)
logger.info("  Batch size = %d", TRAIN_BATCH_SIZE)
logger.info("  Num steps = %d", num_train_optimization_steps)
all_input_ids = torch.tensor([f.input_ids for f in train_features], dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in train_features], dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in train_features], dtype=torch.long)

if OUTPUT_MODE == "classification":
    all_label_ids = torch.tensor([f.label_id for f in train_features], dtype=torch.long)
elif OUTPUT_MODE == "regression":
    all_label_ids = torch.tensor([f.label_id for f in train_features], dtype=torch.float)

train_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids)
train_sampler = RandomSampler(train_data)
train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=TRAIN_BATCH_SIZE)

model.train()
for _ in trange(int(NUM_TRAIN_EPOCHS), desc="Epoch"):
    tr_loss = 0
    nb_tr_examples, nb_tr_steps = 0, 0
    for step, batch in enumerate(tqdm(train_dataloader, desc="Iteration")):
        batch = tuple(t.to(device) for t in batch)
        input_ids, input_mask, segment_ids, label_ids = batch

        # define a new function to compute loss values for both output_modes
        logits = model(input_ids, token_type_ids=segment_ids, attention_mask=input_mask)

        if output_mode == "classification":
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, num_labels), label_ids.view(-1))
        elif output_mode == "regression":
            loss_fct = MSELoss()
            loss = loss_fct(logits.view(-1), label_ids.view(-1))

        if GRADIENT_ACCUMULATION_STEPS > 1:
            loss = loss / GRADIENT_ACCUMULATION_STEPS

        loss.backward()

        tr_loss += loss.item()
        nb_tr_examples += input_ids.size(0)
        nb_tr_steps += 1
        if (step + 1) % GRADIENT_ACCUMULATION_STEPS == 0:
            optimizer.step()
            optimizer.zero_grad()
            global_step += 1

model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self

# If we save using the predefined names, we can load using `from_pretrained`
output_model_file = os.path.join(OUTPUT_DIR, WEIGHTS_NAME)
output_config_file = os.path.join(OUTPUT_DIR, CONFIG_NAME)

torch.save(model_to_save.state_dict(), output_model_file)
model_to_save.config.to_json_file(output_config_file)
<commit_msg>fixing inherited bug<commit_after>'''
Written by Thilina Rajapakse
https://github.com/ThilinaRajapakse/BERT_binary_text_classification/blob/master/converter.py
'''

import torch, sys
import pickle
from math import ceil

from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler,
                              TensorDataset)
from torch.nn import CrossEntropyLoss, MSELoss

from tqdm import tqdm, trange
import os
from pytorch_pretrained_bert import BertTokenizer, BertModel, BertForMaskedLM, BertForSequenceClassification
from pytorch_pretrained_bert.optimization import BertAdam, WarmupLinearSchedule

from multiprocessing import Pool, cpu_count
from tools import *
import convert_examples_to_features

# OPTIONAL: if you want to have more information on what's happening, activate the logger as follows
import logging
logging.basicConfig(level=logging.INFO)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Bert pre-trained model selected in the list: bert-base-uncased,
# bert-large-uncased, bert-base-cased, bert-large-cased, bert-base-multilingual-uncased,
# bert-base-multilingual-cased, bert-base-chinese.
BERT_MODEL = 'bert-base-uncased'

# The name of the task to train.
TASK_NAME = sys.argv[1]

# The input data dir. Should contain the .tsv files (or other data files) for the task.
DATA_DIR = f'data/{TASK_NAME}/'

CACHE_DIR = 'cache/'

# The maximum total input sequence length after WordPiece tokenization.
# Sequences longer than this will be truncated, and sequences shorter than this will be padded.
MAX_SEQ_LENGTH = 128

TRAIN_BATCH_SIZE = 24
EVAL_BATCH_SIZE = 8
LEARNING_RATE = 2e-5
NUM_TRAIN_EPOCHS = 1
RANDOM_SEED = 42
GRADIENT_ACCUMULATION_STEPS = 1
WARMUP_PROPORTION = 0.1

if sys.argv[2] == 'reg':
    OUTPUT_MODE = 'regression'
    # The output directory where the model predictions and checkpoints will be written.
    OUTPUT_DIR = 'outputs/' + TASK_NAME + '_reg/'
else:
    OUTPUT_MODE = 'classification'
    # The output directory where the model predictions and checkpoints will be written.
    OUTPUT_DIR = 'outputs/' + TASK_NAME + '/'

CONFIG_NAME = "config.json"
WEIGHTS_NAME = "pytorch_model.bin"

num_labels = 2

processor = BinaryClassificationProcessor()
train_examples = processor.get_train_examples(DATA_DIR)
train_examples_len = len(train_examples)
print("Train examples count: ", train_examples_len)

with open(OUTPUT_DIR + "train_features.pkl", "rb") as f:
    train_features = pickle.load(f)

train_examples_len = len(train_features)
print("Train features count: ", train_examples_len)

num_train_optimization_steps = ceil(
    train_examples_len / TRAIN_BATCH_SIZE / GRADIENT_ACCUMULATION_STEPS) * NUM_TRAIN_EPOCHS

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = BertForSequenceClassification.from_pretrained('bert-base-uncased', cache_dir='cache', num_labels = num_labels)

model.to(device)

param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
    {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
    {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
    ]

optimizer = BertAdam(optimizer_grouped_parameters,
                     lr=LEARNING_RATE,
                     warmup=WARMUP_PROPORTION,
                     t_total=num_train_optimization_steps)

global_step = 0
nb_tr_steps = 0
tr_loss = 0

logger.info("***** Running training *****")
logger.info("  Num examples = %d", train_examples_len)
logger.info("  Batch size = %d", TRAIN_BATCH_SIZE)
logger.info("  Num steps = %d", num_train_optimization_steps)
all_input_ids = torch.tensor([f.input_ids for f in train_features], dtype=torch.long)
all_input_mask = torch.tensor([f.input_mask for f in train_features], dtype=torch.long)
all_segment_ids = torch.tensor([f.segment_ids for f in train_features], dtype=torch.long)

if OUTPUT_MODE == "classification":
    all_label_ids = torch.tensor([f.label_id for f in train_features], dtype=torch.long)
elif OUTPUT_MODE == "regression":
    all_label_ids = torch.tensor([f.label_id for f in train_features], dtype=torch.float)

train_data = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids)
train_sampler = RandomSampler(train_data)
train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=TRAIN_BATCH_SIZE)

model.train()
for _ in trange(int(NUM_TRAIN_EPOCHS), desc="Epoch"):
    tr_loss = 0
    nb_tr_examples, nb_tr_steps = 0, 0
    for step, batch in enumerate(tqdm(train_dataloader, desc="Iteration")):
        batch = tuple(t.to(device) for t in batch)
        input_ids, input_mask, segment_ids, label_ids = batch

        # define a new function to compute loss values for both output_modes
        logits = model(input_ids, token_type_ids=segment_ids, attention_mask=input_mask)

        if OUTPUT_MODE == "classification":
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, num_labels), label_ids.view(-1))
        elif OUTPUT_MODE == "regression":
            loss_fct = MSELoss()
            loss = loss_fct(logits.view(-1), label_ids.view(-1))

        if GRADIENT_ACCUMULATION_STEPS > 1:
            loss = loss / GRADIENT_ACCUMULATION_STEPS

        loss.backward()

        tr_loss += loss.item()
        nb_tr_examples += input_ids.size(0)
        nb_tr_steps += 1
        if (step + 1) % GRADIENT_ACCUMULATION_STEPS == 0:
            optimizer.step()
            optimizer.zero_grad()
            global_step += 1

model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self

# If we save using the predefined names, we can load using `from_pretrained`
output_model_file = os.path.join(OUTPUT_DIR, WEIGHTS_NAME)
output_config_file = os.path.join(OUTPUT_DIR, CONFIG_NAME)

torch.save(model_to_save.state_dict(), output_model_file)
model_to_save.config.to_json_file(output_config_file)
<|endoftext|>


<commit_before>#!/usr/bin/python3

import logging

import pyawskit.common

logging.basicConfig()
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

"""
This script mounts all the local disks as individuals

TODO:
- make this run in parallel on multiple cores and enable the user to choose (via
command line option) whether to run this multi-core or not.
"""

def main():
    disks = pyawskit.common.get_disks()
    for disk in disks:
        folder = "/mnt/{}".format(disk)
        pyawskit.common.erase_partition_table(disk=disk)
        pyawskit.common.format_device(disk=disk)
        pyawskit.common.mount_disk(disk=disk, folder=folder)

if __name__ == "__main__":
    main()
<commit_msg>some more TODO items<commit_after>#!/usr/bin/python3

import logging

import pyawskit.common

logging.basicConfig()
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

"""
This script mounts all the local disks as individuals

TODO:
- make this run in parallel on multiple cores and enable the user to choose (via
command line option) whether to run this multi-core or not.
"""

def main():
    # TODO: check that we are running as root
    # if not then recommend sudo(1) and exit
    disks = pyawskit.common.get_disks()
    for disk in disks:
        folder = "/mnt/{}".format(disk)
        pyawskit.common.erase_partition_table(disk=disk)
        pyawskit.common.format_device(disk=disk)
        pyawskit.common.mount_disk(disk=disk, folder=folder)

if __name__ == "__main__":
    main()
<|endoftext|>