text:
import re
import os
import glob
import numpy as np
import random
import random
import pandas as pd
import json
from torch.utils.data import Dataset
import logging
from tqdm import tqdm
#import spacy
#nlp = spacy.load("en_core_web_sm")
LOGGER = logging.getLogger(__name__)

def delete_tokenmark_input(input_ids,tokenizer):
    input_id_new=[]
    del_num=0
    token_pos_start_id=[tokenizer.encode('[',add_special_tokens=False)[0],tokenizer.encode(' [',add_special_tokens=False)[0]]
    token_pos_end_id=[tokenizer.encode(']',add_special_tokens=False)[0],tokenizer.encode(' ]',add_special_tokens=False)[0]]
    token_pos_start_end_id=set(token_pos_start_id+token_pos_end_id)
    for i,input_i in enumerate(input_ids):
        if input_i not in token_pos_start_end_id:
            input_id_new.append(input_i)
        else:
            del_num+=1
    input_id_new+=del_num*[tokenizer.pad_token_id]
    return input_id_new

def delete_tokenmarker_am(input_ids,tokenizer):
    am_new=[]
    for i in input_ids:
        if i==tokenizer.pad_token_id:
            am_new.append(0)
        else:
            am_new.append(1)
    return am_new

def find_token_id(input_id,tokenizer):
    token_pos_start_id=set([tokenizer.encode('[',add_special_tokens=False)[0],tokenizer.encode(' [',add_special_tokens=False)[0]])
    token_pos_end_id=set([tokenizer.encode(']',add_special_tokens=False)[0],tokenizer.encode(' ]',add_special_tokens=False)[0]])

    token_ids=[]
    for i,input_i in enumerate(input_id):
        input_i=int(input_i)
        if i==len(input_id)-1: # the last token
            continue
        if input_i in [tokenizer.mask_token_id,tokenizer.cls_token_id,tokenizer.pad_token_id]:
            continue
        if input_i in token_pos_start_id:
            token_ids.append(i+1)
            # logger.info("first word",token_ids)
        elif input_i in token_pos_end_id:
            token_ids.append(i)
    try:
        assert len(token_ids)==2
    except AssertionError as e:
        print ('Warning, token id alter is not length 2')
        print (input_id)
        print (tokenizer.convert_ids_to_tokens(input_id))
        print (token_pos_start_id)
        print (token_pos_end_id)
        print (token_ids)
        sys.exit(1)

    try:
       assert token_ids[1]!=token_ids[0]
    except AssertionError as e:
        print ('token marker star == end')
        print (input_id)
        print (tokenizer.convert_ids_to_tokens(input_id))
        print (token_ids)
        sys.exit(1)
    token_ids[1]=token_ids[1]-1
    token_ids[0]=token_ids[0]-1
    return token_ids

def delete_tokenmaker_tokentypeids(input_ids,tokenizer):
    tokentype_ids=[]
    item=0
    for i in input_ids:

        if i==tokenizer.pad_token_id:
            tokentype_ids.append(0)

        elif i==tokenizer.sep_token_id:
            tokentype_ids.append(item)
            item=1
        else:
            tokentype_ids.append(item)
    return tokentype_ids

class QueryDataset_custom(Dataset):

    def __init__(self, data_dir,
                load_full_sentence=False,
                filter_duplicate=False
        ):
        """
        Parameters
        ----------
        data_dir : str
            a path of data
        filter_duplicate : bool
            filter duplicate queries
        draft : bool
            use subset of queries for debugging (default False)
        """
        LOGGER.info("QueryDataset! data_dir={} filter_duplicate={}".format(
            data_dir, filter_duplicate
        ))

        self.data = self.load_data(
            data_dir=data_dir,
            filter_duplicate=filter_duplicate
        )

    def load_data(self, data_dir, filter_duplicate):
        """
        Parameters
        ----------
        data_dir : str
            a path of data
        filter_composite : bool
            filter composite mentions
        filter_duplicate : bool
            filter duplicate queries

        Returns
        -------
        data : np.array
            mention, cui pairs
        """
        data = []

        with open(data_dir, "r") as f:
            lines = f.readlines()

        for line in lines:
            line = line.rstrip("\n")
            _id, mention = line.split("||")

            data.append((mention, _id))

        if filter_duplicate:
            data = list(dict.fromkeys(data))

        # return np.array data
        data = np.array(data, dtype=object)

        return data

class QueryDataset_pretraining(Dataset):

    def __init__(self, data_dir,
                filter_duplicate=False
        ):
        """
        Parameters
        ----------
        data_dir : str
            a path of data
        filter_duplicate : bool
            filter duplicate queries
        draft : bool
            use subset of queries for debugging (default False)
        """
        LOGGER.info("QueryDataset! data_dir={} filter_duplicate={}".format(
            data_dir,filter_duplicate
        ))

        self.data = self.load_data(
            data_dir=data_dir,
            filter_duplicate=filter_duplicate
        )

    def load_data(self, data_dir, filter_duplicate):
        """
        Parameters
        ----------
        data_dir : str
            a path of data
        filter_composite : bool
            filter composite mentions
        filter_duplicate : bool
            filter duplicate queries

        Returns
        -------
        data : np.array
            mention, cui pairs
        """
        data = []

        #concept_files = glob.glob(os.path.join(data_dir, "*.concept"))
        with open(data_dir, "r") as f:
            lines = f.readlines()

        for row in lines:
            row = row.rstrip("\n")
            snomed_id, mention = row.split("||")
            data.append((mention, snomed_id))

        if filter_duplicate:
            data = list(dict.fromkeys(data))

        # return np.array data
        data = np.array(data)

        return data

class QueryDataset(Dataset):

    def __init__(self, data_dir,
                filter_composite=False,
                filter_duplicate=False
        ):
        """
        Parameters
        ----------
        data_dir : str
            a path of data
        filter_composite : bool
            filter composite mentions
        filter_duplicate : bool
            filter duplicate queries
        draft : bool
            use subset of queries for debugging (default False)
        """
        LOGGER.info("QueryDataset! data_dir={} filter_composite={} filter_duplicate={}".format(
            data_dir, filter_composite, filter_duplicate
        ))

        self.data = self.load_data(
            data_dir=data_dir,
            filter_composite=filter_composite,
            filter_duplicate=filter_duplicate
        )

    def load_data(self, data_dir, filter_composite, filter_duplicate):
        """
        Parameters
        ----------
        data_dir : str
            a path of data
        filter_composite : bool
            filter composite mentions
        filter_duplicate : bool
            filter duplicate queries

        Returns
        -------
        data : np.array
            mention, cui pairs
        """
        data = []

        #concept_files = glob.glob(os.path.join(data_dir, "*.txt"))
        file_types = ("*.concept", "*.txt")
        concept_files = []
        for ft in file_types:
            concept_files.extend(glob.glob(os.path.join(data_dir, ft)))

        for concept_file in tqdm(concept_files):
            with open(concept_file, "r", encoding='utf-8') as f:
                concepts = f.readlines()

            for concept in concepts:
                #print (concept)
                concept = concept.split("||")
                #if len(concept) !=5: continue
                mention = concept[3].strip().lower()
                cui = concept[4].strip()
                if cui.lower() =="cui-less": continue
                is_composite = (cui.replace("+","|").count("|") > 0)

                if filter_composite and is_composite:
                    continue
                else:
                    data.append((mention,cui))

        if filter_duplicate:
            data = list(dict.fromkeys(data))

        # return np.array data
        data = np.array(data)

        return data

class DictionaryDataset():
    """
    A class used to load dictionary data
    """
    def __init__(self, dictionary_path):
        """
        Parameters
        ----------
        dictionary_path : str
            The path of the dictionary
        draft : bool
            use only small subset
        """
        LOGGER.info("DictionaryDataset! dictionary_path={}".format(
            dictionary_path
        ))
        self.data = self.load_data(dictionary_path)

    def load_data(self, dictionary_path):
        name_cui_map = {}
        data = []
        with open(dictionary_path, mode='r', encoding='utf-8') as f:
            lines = f.readlines()
            for line in tqdm(lines):
                line = line.strip()
                if line == "": continue
                cui, name = line.split("||")
                name = name.lower()
                if cui.lower() == "cui-less": continue
                data.append((name,cui))

        #LOGGER.info("concerting loaded dictionary data to numpy array...")
        #data = np.array(data)
        return data

def erase_and_mask_token(s, erase_len=3):
    for i,token in enumerate(s):
        if token=='[':
           token_start=i
        elif token==']':
            token_end=i+1
    token_left=s[:token_start]
    token_right=s[token_end:]
    token=s[token_start:token_end]
    if random.sample([1,2],1)[0]==1:
        token_left=reorder(token_left,erase_len)
    else:
        token_right=reorder(token_right,erase_len)
    return token_left+' '+token+ ' '+token_right

def reorder(s,shuffle_len):
    wlist=s.split()
    if len(wlist) <= shuffle_len: return s
    ind = np.random.randint(len(wlist)-shuffle_len)
    select_shuffle=wlist[ind:ind+shuffle_len]
    random.shuffle(select_shuffle)
    wlist=wlist[:ind]+select_shuffle+wlist[ind+shuffle_len:]
    return ' '.join(wlist)

def erase_and_mask(s, erase_len=3):
    if len(s) <= erase_len: return s
    if len(s) < 5: return s # if too short, no augmentation
    ind = np.random.randint(len(s)-erase_len)
    left, right = s.split(s[ind:ind+erase_len], 1)
    return " ".join([left, "[MASK]", right])

class MetricLearningDataset_pairwise(Dataset):
    """
    Candidate Dataset for:
        query_tokens, candidate_tokens, label
    """
    def __init__(self, path, tokenizer, random_erase=0): #d_ratio, s_score_matrix, s_candidate_idxs):
        with open(path, 'r') as f:
            lines = f.readlines()
        self.query_ids = []
        self.query_names = []
        for line in lines:
            line = line.rstrip("\n")
            query_id, name1, name2 = line.split("||")
            self.query_ids.append(query_id)
            self.query_names.append((name1, name2))
        self.tokenizer = tokenizer
        self.query_id_2_index_id = {k: v for v, k in enumerate(list(set(self.query_ids)))}
        self.random_erase = random_erase

    def __getitem__(self, query_idx):

        query_name1 = self.query_names[query_idx][0]
        query_name2 = self.query_names[query_idx][1]
        if self.random_erase != 0:
            query_name2 = erase_and_mask_token(query_name2, erase_len=int(self.random_erase))
            # print (query_name2)
        query_id_orig = self.query_ids[query_idx]
        query_id = int(self.query_id_2_index_id[query_id_orig])

        return query_name1, query_name2, query_id, query_id_orig

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

max_stars_repo_name: cambridgeltl/MirrorWiC

max_issues_repo_name: cambridgeltl/MirrorWiC

max_forks_repo_name: cambridgeltl/MirrorWiC


text:
from neat.species import Species
from neat.genome import Genome
from neat.genes import ConnectionGene, NodeGene, NodeType
from neat.activation_functions import ActivationFunction
from neat.neural_network import Neuron, Connection, NeuralNetwork

import json
import copy
import random
import math
import numpy as np
from dataclasses import dataclass

@dataclass
class PopulationParams:
	population_max: int = 0
	generations: int = 0
	distance_coeff_1: float = 0.0
	distance_coeff_2: float = 0.0
	distance_coeff_3: float = 0.0
	distance_threshold: float = 0.0
	small_genome_coeff: float = 0.0
	no_crossover_offspring: float = 0.0
	survival_selection: bool = False
	allow_clones: bool = False
	survival_threshold: float = 0.0
	elite_offspring_param: float = 0.0
	min_species: int = 0
	max_species: int = 0
	stagnation_purge: bool = False

@dataclass
class PopulationProb:
	# Interspecies probability of mate
	interspecies_mating: float = 0.0

	# Small population mutation probabilities
	sp_new_node: float = 0.0
	sp_new_connection: float = 0.0

	# Large population mutation probabilities
	lp_new_node: float = 0.0
	lp_new_connection: float = 0.0

	# Probability to change enable status of connections
	mutate_connection_status: float = 0.0

	# Mutation probabilities
	mutation_weight: float = 0.0
	mutate_activation: float = 0.0

class NodeHistoryStruct:

	def __init__(self, incoming_connection_id, outgoing_connection_id, node_connected):
		self.incoming_connection_id = incoming_connection_id
		self.outgoing_connection_id = outgoing_connection_id
		self.node_connected = node_connected

class Population:

	def __init__(self):
		self.params = PopulationParams()
		self.prob = PopulationProb()

		self.current_generation = 0
		self.champion_fitness = 0.0
		self.champion_genome = None
		self.best_historical_fitness = 0.0

		self.global_innovation_count = 0
		self.global_node_count = 0
		self.global_species_count = -1

		self.organisms = []
		self.offspring_organisms = []
		self.species = {}

		self.conn_innovation_history = {}
		self.node_history = {}
		self.species_history = {}

		self.distance_threshold_var = 0.0
		self.speciation_adjust_start = False

		self.activation_set = ActivationFunction()
		# Must be set True to operate
		self.configurated = False

	def config_population(self, config):
		self.params.population_max = config['populationMax']
		self.params.generations = config['generations']
		self.params.distance_coeff_1 = config['distanceCoeff1']
		self.params.distance_coeff_2 = config['distanceCoeff2']
		self.params.distance_coeff_3 = config['distanceCoeff3']
		self.params.distance_threshold = config['distanceThreshold']
		self.params.small_genome_coeff = config['smallGenomeCoeff']
		self.params.no_crossover_offspring = config['percentageOffspringWithoutCrossover']

		self.params.survival_selection = config['survivalSelection']
		self.params.allow_clones = config['allowClones']
		self.params.survival_threshold = config['survivalThreshold']
		self.params.elite_offspring_param = config['eliteOffspringParam']

		self.params.min_species = config['adjustMinSpecies']
		self.params.max_species = config['adjustMaxSpecies']

		self.params.stagnation_purge = config['stagnationPurge']

		# Interspecies probability of mate
		self.prob.interspecies_mating = config['probInterspeciesMating']

		# Small population mutation probabilities
		self.prob.sp_new_node = config['sp_probAddingNewNode']
		self.prob.sp_new_connection = config['sp_probAddingNewConnection']

		# Large population mutation probabilities
		self.prob.lp_new_node = config['lp_probAddingNewNode']
		self.prob.lp_new_connection = config['lp_probAddingNewConnection']

		# Probability to change enable status of connections
		self.prob.mutate_connection_status = config['probMutateEnableConnection']

		# Mutation probabilities
		self.prob.mutation_weight = config['probChangeWeight']
		self.prob.mutate_activation = config['probChangeNodeFunction']

		self.configurated = True

	def get_new_innovation(self):
		self.global_innovation_count += 1
		return self.global_innovation_count

	def get_new_node_id(self):
		self.global_node_count += 1
		return self.global_node_count

	def get_new_species_id(self):
		self.global_species_count += 1
		return self.global_species_count

	def start_population(self, genome):
		self.seed_genome = copy.deepcopy(genome)

		for node in genome.node_list:
			if self.global_node_count < node.gene_id:
				self.global_node_count = node.gene_id

		for connection in genome.connection_list:
			self.conn_innovation_history[(connection.incoming, connection.outgoing)] = connection.innovation
			if self.global_innovation_count < connection.innovation:
				self.global_innovation_count = connection.innovation

		self.organisms = [copy.deepcopy(genome) for i in range(self.params.population_max)]

		for org in self.organisms:
			org.randomize_weights()

		self.champion_fitness = 0.0
		self.speciate()

	def restart_population(self):
		self.current_generation = 0
		self.champion_fitness = 0.0
		self.best_historical_fitness = 0.0

		self.global_innovation_count = 0
		self.global_node_count = 0
		self.global_species_count = -1

		self.organisms = []
		self.offspring_organisms = []
		self.species = {}

		self.conn_innovation_history = {}
		self.node_history = {}

		self.distance_threshold_var = 0.0
		self.speciation_adjust_start = False

		for node in self.seed_genome.node_list:
			if self.global_node_count < node.gene_id:
				self.global_node_count = node.gene_id

		for connection in self.seed_genome.connection_list:
			self.conn_innovation_history[(connection.incoming, connection.outgoing)] = connection.innovation
			if self.global_innovation_count < connection.innovation:
				self.global_innovation_count = connection.innovation

		self.organisms = [copy.deepcopy(self.seed_genome) for i in range(self.params.population_max)]

		for org in self.organisms:
			org.randomize_weights()

		self.speciate()

	def extinction(self):
		for org in self.organisms:
			org.randomize_weights()

		self.speciate()

	def compatibility(self, org_A, org_B):
		excess = 0
		disjoint = 0
		var_weight = 0

		itr_A = 0
		itr_B = 0

		conn_A = org_A.connection_list[itr_A]
		conn_B = org_B.connection_list[itr_B]

		larger_genome = max(len(org_A.connection_list), len(org_B.connection_list))

		while True:
			if conn_A.innovation > conn_B.innovation:
				disjoint += 1
				itr_B += 1
			elif conn_A.innovation == conn_B.innovation:
				var_weight += abs(conn_A.weight - conn_B.weight)
				itr_A += 1
				itr_B += 1
			else:
				disjoint += 1
				itr_B += 1

			if itr_A >= len(org_A.connection_list):
				while itr_B < len(org_B.connection_list):
					excess += 1
					itr_B += 1
				break

			if itr_B >= len(org_B.connection_list):
				while itr_A < len(org_A.connection_list):
					excess += 1
					itr_A += 1
				break

			conn_A = org_A.connection_list[itr_A]
			conn_B = org_B.connection_list[itr_B]

		if larger_genome > self.params.small_genome_coeff:
			divisor = larger_genome
			small_genomes_buff = 1.0 / larger_genome
		else:
			divisor = 1
			if self.params.distance_coeff_3 < 1.0:
				small_genomes_buff = 1.0 / self.params.distance_coeff_3
			else:
				small_genomes_buff = 1.0

		compatibility = self.params.distance_coeff_1 * excess / divisor + self.params.distance_coeff_2 * disjoint / divisor + self.params.distance_coeff_3 * var_weight * small_genomes_buff
		return compatibility

	def sort_organisms(self):
		self.organisms.sort(key=lambda x: x.fitness, reverse=True)

	# From Stanley code
	def adjust_speciate_threshold(self):
		re_speciate = False

		if len(self.species) < self.params.min_species:
			if self.speciation_adjust_start == False:
				self.speciation_adjust_start = True
				self.distance_threshold_var = -max(0.1, self.params.distance_threshold * 0.01)
			else:
				if self.distance_threshold_var < 0.0:
					self.distance_threshold_var *= 1.05
				else:
					self.distance_threshold_var *= -0.5

			self.params.distance_threshold += self.distance_threshold_var
			self.params.distance_threshold = max(0.1, self.params.distance_threshold)

			re_speciate = False # True

		elif len(self.species) > self.params.max_species:
			if self.speciation_adjust_start == False:
				self.speciation_adjust_start = True
				self.distance_threshold_var = max(0.1, self.params.distance_threshold * 0.01)
			else:
				if self.distance_threshold_var < 0.0:
					self.distance_threshold_var *= -0.5
				else:
					self.distance_threshold_var *= 1.05

			self.params.distance_threshold += self.distance_threshold_var

			re_speciate = False # True

		else:
			self.speciation_adjust_start = False

		if re_speciate is True:
			self.species = {}
			self.speciate()

	def crossover(self, org_A, org_B):
		# Iterator used to navigate through connection_list and node_list independently for each organism. TODO: use iter()
		itr_A = 0
		itr_B = 0

		# connection_gene's and node_gene's to be selected as a result of crossover
		conn_A = org_A.connection_list[itr_A]
		conn_B = org_A.connection_list[itr_B]

		node_A = org_A.node_list[itr_A]
		node_B = org_B.node_list[itr_B]

		# Output genome of the crossover operation. #TODO: enhance genome copy
		new_organism = Genome(num_layers=org_A.num_layers, phenotype=org_A.phenotype, weights_range=self.seed_genome.weights_range)
		new_organism.parent_species = org_A.parent_species

		# Start iterate over genomes to explore the connection_list
		while True:
			if conn_A.innovation > conn_B.innovation:
				new_organism.add_connection(conn_B)
				itr_B += 1
			elif conn_A.innovation == conn_B.innovation:
				if conn_A.enable == True and conn_B.enable == True:
					new_organism.add_connection(random.choice([conn_A, conn_B]))
				elif conn_A.enable == True:
					new_organism.add_connection(conn_A)
				else:
					new_organism.add_connection(conn_B)

				itr_A += 1
				itr_B += 1
			else:
				new_organism.add_connection(conn_A)
				itr_A += 1

			# If iterator is bigger than any of the genome, it means excess
			if itr_A >= len(org_A.connection_list):
				while itr_B < len(org_B.connection_list):
					new_organism.add_connection(org_B.connection_list[itr_B])
					itr_B += 1
				break

			if itr_B >= len(org_B.connection_list):
				while itr_A < len(org_A.connection_list):
					new_organism.add_connection(org_A.connection_list[itr_A])
					itr_A += 1
				break

			conn_A = org_A.connection_list[itr_A]
			conn_B = org_B.connection_list[itr_B]

		itr_A = 0
		itr_B = 0

		# Now iterate over node_list
		while True:
			if node_A.gene_id > node_B.gene_id:
				new_organism.add_node(node_B)
				itr_B += 1
			elif node_A.gene_id == node_B.gene_id:
				new_organism.add_node(random.choice([node_A, node_B]))
				itr_A += 1
				itr_B += 1
			else:
				new_organism.add_node(node_A)
				itr_A += 1

			# If iterator is bigger than any of the genome, it means excess
			if itr_A >= len(org_A.node_list):
				while itr_B < len(org_B.node_list):
					new_organism.add_node(org_B.node_list[itr_B])
					itr_B += 1
				break

			if itr_B >= len(org_B.node_list):
				while itr_A < len(org_A.node_list):
					new_organism.add_node(org_A.node_list[itr_A])
					itr_A += 1
				break

			node_A = org_A.node_list[itr_A]
			node_B = org_B.node_list[itr_B]

		return new_organism

	def mutate_add_node(self, organism):
		if len(organism.connection_list) == 0:
			return

		# Select connection to replace
		connection_replace = random.choice(organism.connection_list)

		if connection_replace.innovation < 0 or connection_replace.enable == False:
			return

		# Check if innovation already mutated
		if connection_replace.innovation in self.node_history:
			new_node = self.node_history[connection_replace.innovation].node_connected
			new_node.randomize_function(self.activation_set)

			# Load incoming and outgoing innovation from the node_history dict
			incoming_connection_id = self.node_history[connection_replace.innovation].incoming_connection_id
			outgoing_connection_id = self.node_history[connection_replace.innovation].outgoing_connection_id

			# Create new connections
			incoming_connection = ConnectionGene(incoming_connection_id, connection_replace.incoming, new_node.gene_id, 1.0, True, connection_replace.source_layer, new_node.layer)
			outgoing_connection = ConnectionGene(outgoing_connection_id, new_node.gene_id, connection_replace.outgoing, 0.0, True, new_node.layer, connection_replace.target_layer)
			outgoing_connection.randomize_weight(organism.weights_range[0], organism.weights_range[1])
		else:
			new_node_layer = math.floor(abs(connection_replace.target_layer - connection_replace.source_layer) / 2) + min(connection_replace.source_layer, connection_replace.target_layer)
			if connection_replace.target_layer < new_node_layer or new_node_layer == connection_replace.source_layer or new_node_layer == connection_replace.target_layer:
				return

			new_node = NodeGene(self.get_new_node_id(), NodeType.HIDDEN, self.activation_set.get_random_function(), new_node_layer)

			# Create new connections and increment innovation
			incoming_connection = ConnectionGene(self.get_new_innovation(), connection_replace.incoming, new_node.gene_id, 1.0, True, connection_replace.source_layer, new_node_layer)
			outgoing_connection = ConnectionGene(self.get_new_innovation(), new_node.gene_id, connection_replace.outgoing, 0.0, True, new_node_layer, connection_replace.target_layer)
			outgoing_connection.randomize_weight(organism.weights_range[0], organism.weights_range[1])

			# Register new additions to node_history
			self.node_history[connection_replace.innovation] = NodeHistoryStruct(incoming_connection.innovation, outgoing_connection.innovation, new_node)

		# Disable replaced connection
		connection_replace.enable = False

		organism.add_node(new_node)
		organism.add_connection(incoming_connection)
		organism.add_connection(outgoing_connection)

	def mutate_add_connection(self, organism):
		#TODO: add conditions to non-recursive and feed-forward
		input_candidate = random.choice(organism.node_list)
		source_layer_candidate = input_candidate.layer

		# Search for outgoing nodes
		output_candidates = []
		for node in organism.node_list:
			if node.node_type != NodeType.INPUT and node.node_type != NodeType.BIAS:
				output_candidates.append(node)

		if len(output_candidates) < 1:
			return

		# Select a node and verify that is not the same as input
		output_candidate = random.choice(output_candidates)
		target_layer_candidate = output_candidate.layer
		if input_candidate == output_candidate:
			return

		# Check for loops in the net
		#if self.check_loops(organism, input_candidate.gene_id, output_candidate.gene_id) is True:
		#	return

		# Make sure the connection doesn't exist already
		node_pair = (input_candidate.gene_id, output_candidate.gene_id)
		connection_index = -1
		for index, conn in enumerate(organism.connection_list):
			if conn.incoming == node_pair[0] and conn.outgoing == node_pair[1]:
				connection_index = index
				break

		if connection_index != -1:
			new_connection = ConnectionGene(incoming=node_pair[0], outgoing=node_pair[1], enable=True, source_layer=source_layer_candidate, target_layer=target_layer_candidate)

			if node_pair in self.conn_innovation_history:
				new_connection.innovation = self.conn_innovation_history[node_pair]
				new_connection.randomize_weight(organism.weights_range[0], organism.weights_range[1])
			else:
				# Increment global innovation before creating new ConnectionGene. #TODO make method that do this
				new_connection.innovation = self.get_new_innovation()
				self.conn_innovation_history[node_pair] = new_connection.innovation

			if new_connection.innovation > self.global_innovation_count:
				raise NameError('Global innovation mismatched')

			organism.add_connection(new_connection)
		else:
			organism.connection_list[connection_index].randomize_weight(organism.weights_range[0], organism.weights_range[1])

	def mutate_connection_weight(self, organism):
		for conn in organism.connection_list:
			if random.uniform(0, 1) < self.prob.mutation_weight:
				if not conn.mutated:
					conn.randomize_weight(organism.weights_range[0], organism.weights_range[1])

	def mutate_node_functions(self, organism):
		for node in organism.node_list:
			if random.uniform(0, 1) < self.prob.mutate_activation and node.node_type != NodeType.INPUT:
				node.randomize_function(self.activation_set)

	def mutation(self, org):
		org.recompute_complexity()
		if org.complexity > self.params.small_genome_coeff:
			probabilities_set = {self.prob.lp_new_node: self.mutate_add_node(org),
								self.prob.lp_new_connection: self.mutate_add_connection(org),
								self.prob.mutate_activation: self.mutate_node_functions(org),
								self.prob.mutation_weight: self.mutate_connection_weight(org)}
		else:
			probabilities_set = {self.prob.sp_new_node: self.mutate_add_node(org),
								self.prob.sp_new_connection: self.mutate_add_connection(org),
								self.prob.mutate_activation: self.mutate_node_functions(org),
								self.prob.mutation_weight: self.mutate_connection_weight(org)}

		probabilities_sum = np.fromiter(probabilities_set.keys(), dtype=float).sum()
		probabilities_accum = 0.0

		prob_eval = random.uniform(0, probabilities_sum)

		for prob in probabilities_set:
			probabilities_accum += prob
			if prob_eval < probabilities_accum and prob != 0.0:
				probabilities_set[prob]
				break

	def reproduce(self):
		self.offspring_organisms = []

		for sp in self.species.values():
			if sp.extinct == True:
				continue

			offspring_amount = math.floor(sp.offspring)

			if offspring_amount == 0:
				continue

			elite_offspring = round(len(sp.organisms) * self.params.elite_offspring_param)
			elite_count = 0

			if elite_offspring == 0:
				elite_offspring = 1

			while offspring_amount > 0:
				if elite_count < elite_offspring:
					son = copy.deepcopy(sp.best_organism)
					elite_count += 1
				else:
					if random.uniform(0, 1) < self.params.no_crossover_offspring:
						son = copy.deepcopy(self.get_random_organism(sp))
					else:
						random_mother = self.get_random_organism(sp)
						if len(self.species) > 1 and random.uniform(0, 1) < self.prob.interspecies_mating:
							while True:
								random_species = random.choice(list(self.species.values()))
								if(sp.birth != random_species.birth):
									break
							random_father = self.get_random_organism(sp)
							son = self.crossover(random_mother, random_father)
						else:
							if len(sp.organisms) == 1:
								son = copy.deepcopy(random_mother)
							else:
								while True:
									random_father = self.get_random_organism(sp)
									if(random_mother is not random_father or self.params.allow_clones):
										break
								son = self.crossover(random_mother, random_father)

					self.mutation(son)

				offspring_amount -= 1
				self.offspring_organisms.append(son)

		# Make sure the population is the correct size
		while len(self.offspring_organisms) < self.params.population_max:
			random_species = random.choice(list(self.species.values()))

			random_org = copy.deepcopy(random_species.best_organism)
			random_org.randomize_weights()

			self.offspring_organisms.append(random_org)

		self.organisms = self.offspring_organisms

	def speciate(self):
		speciated = []

		for sp in self.species.values():
			sp.organisms = []

			if sp.extinct == True:
				continue

			distances = []
			for org in self.organisms:
				distances.append(self.compatibility(org, sp.representant))

			if np.min(distances) >= self.params.distance_threshold:
				sp.extinct = True
				continue

			idx = np.argmin(distances)
			sp.representant = self.organisms[idx]
			sp.organisms.append(self.organisms[idx])

			self.organisms[idx].parent_species = sp.birth
			speciated.append(self.organisms.pop(idx))

		while self.organisms:
			compatible_species = False

			if len(self.species) == 0:
				new_species = Species()
				new_species.birth = self.get_new_species_id()
				new_species.organisms.append(self.organisms[0])

				new_species.best_organism = self.organisms[0]
				new_species.representant = self.organisms[0]

				self.organisms[0].parent_species = new_species.birth # TODO: a method to assign and post-update new count

				self.species[new_species.birth] = new_species

				speciated.append(self.organisms.pop(0))

			else:
				"""
				if org.parent_species != -1 and org.parent_species in self.species:
					compatibility = self.compatibility(org, self.species[org.parent_species].best_organism)
					print('Compability vs threshold: {:.4f} - {:.4f}'.format(compatibility, self.params.distance_threshold))
					if compatibility < self.params.distance_threshold:
						self.species[org.parent_species].organisms.append(org)
						continue

					setscore
				"""

				distances = []
				for key, sp in self.species.items():
					if sp.extinct == True:
						continue
					distances.append([key, self.compatibility(self.organisms[0], sp.representant)])
					#print('Compability vs threshold: {:.4f} - {:.4f}'.format(compatibility, self.params.distance_threshold))

				if len(distances) > 0:
					min_idx, min_dist = min(distances, key=lambda x: x[1])

					if min_dist < self.params.distance_threshold:
						compatible_species = True
						self.species[min_idx].organisms.append(self.organisms[0])

						self.organisms[0].parent_species = min_idx
						speciated.append(self.organisms.pop(0))

				if compatible_species is False:
					new_species = Species()
					new_species.birth = self.get_new_species_id()
					new_species.organisms.append(self.organisms[0])

					new_species.best_organism = self.organisms[0]
					new_species.representant = self.organisms[0]

					self.organisms[0].parent_species = new_species.birth # TODO: a method to assign and post-update new count

					self.species[new_species.birth] = new_species

					speciated.append(self.organisms.pop(0))

		self.organisms = speciated

	def check_loops(self, organism, inspect, itr):
		for conn in organism.connection_list:
			if conn.incoming == itr:
				if conn.outgoing == inspect:
					return True
				else:
					self.check_loops(organism, inspect, conn.outgoing)

		return False

	#def sort_species_by_fitness(self):
	#	self.species.sort(key=lambda x: x.best_fitness, reverse=True)

	def compute_offspring(self):
		pop_shared_fitness = 0.0
		pop_avg_shared_fitness = 0.0

		for sp in self.species.values():
			sp.age += 1
			sp.offspring = 0.0
			sp.avg_fitness = 0.0

			sp.update_champion()
			if sp.age - sp.last_improvement_age > 8: # TESTING
				sp.extinct = True

			if sp.best_organism.fitness >= self.champion_fitness:
				self.champion_fitness = sp.best_organism.fitness
				self.champion_genome = copy.deepcopy(sp.best_organism)

			for org in sp.organisms:
				org.shared_fitness = org.fitness / len(sp.organisms)
				sp.avg_fitness += org.shared_fitness

			pop_shared_fitness += sp.avg_fitness

		pop_avg_shared_fitness = pop_shared_fitness / self.params.population_max

		for sp in self.species.values():
			#if sp.extinct == True:
			#	continue
			if pop_avg_shared_fitness == 0 or sp.extinct == True:
				sp.offspring = 0.0
			else:
				sp.offspring = sp.avg_fitness / pop_avg_shared_fitness

	def remove_empty_species(self):
		species_to_remove = []
		for sp_id in self.species:
			if len(self.species[sp_id].organisms) == 0 or self.species[sp_id].extinct == True:
				species_to_remove.append(sp_id)
				continue

		for sp in species_to_remove:
			sp_del = self.species.pop(sp)
			del sp_del

	def get_random_organism(self, species):
		if len(species.organisms) == 0:
			return None

		survival_cap = math.floor(self.params.survival_threshold * len(species.organisms))

		if survival_cap == 0:
			return species.organisms[0]
		else:
			return random.choice(species.organisms[:survival_cap])

ext: py

max_stars_repo_name: pabloreyesrobles/py-hyperneat

max_issues_repo_name: pabloreyesrobles/py-hyperneat

max_forks_repo_name: pabloreyesrobles/py-hyperneat

reason: import numpy


text:
import itertools
import numpy as np
import EDGE as edge
from astropy.io import ascii
import pdb

'''
jobmaker.py

PURPOSE:
    Script that uses job_file_create to produce a file with all the parameters in the grid and creates the job files
    themselves.

HOW TO USE THIS SCRIPT:
    1) Change the gridpath to be where you want the jobfiles to be placed. NOTE: This should also be the location of the jobsample file.

    2) Change the paramfiletag to an insightful name that will help you identify the grid of models you ran three months from now.

    3) Change the parameters in the brackets below to whatever set of parameters you want to run.
       The parameters should be surrounded by []'s and be separated by a commas (A list)

    4) Change the labelname to the name of the object, e.g. 'gmauriga'

    5) Run the script

OPTIONAL PARAMETERS:
    'amaxw' and 'mdotstar': Generally these take the value of 'amaxs' and 'mdot' respectively, but it is possible that you may want to change them.
    To do this, change them like you would any other parameter. If you DO want to keep them the same as 'amaxs' and 'mdot', then set the list equal to None

    fill: This will change the zero padding for the job files.

    iwall: [Boolean] Turns on/off all the disk components and just runs the wall. Generally set to 'False'
           Note: This is not looped over (i.e., all the models with this set to 'True' will only be walls).

NOTES:
    'amaxs', 'amaxw', and lamaxb only accept certain values. 'amaxw' accepts the same values as 'amaxs'
    Here are the possible values:
        amaxs/amaxw: ['0.05', '0.1', '0.25', '0.5', '0.75', '1.0', '1.25',
        '1.5', '1.75', '2.0', '2.25', '2.5', '3.0', '4.0', '5.0', '10.0',
        '100.0']

        lamaxb: ['100','200','300','400',500', '600', '700', '800',' 900',
        '1mm', '2mm', '3mm', '4mm', '5mm', '6mm', '7mm', '8mm', '9mm', '1cm',
        '1p1cm', '1p2cm', '1p3cm', '1p4cm', '1p5cm', '1p6cm', '1p7cm', '1p8cm',
        '1p9cm', '2cm', '2p1cm', '2p2cm', '2p3cm', '2p4cm', '2p5cm']

MODIFICATION HISTORY
    Jun 16th, 2017 - Updated for mdotstar + amaxw by Connor Robinson
    Apr 29th, 2016 - Written by Connor Robinson
'''

#Where you want the parameter file and the jobfiles to be placed
#Also must be the location of the sample job file
gridpath = '/Users/Connor/Desktop/Research/diad/test/'

#Where you will be running the jobs on the cluster
clusterpath = '/projectnb/bu-disks/connorr/test/'

#Tag that you can add to make the parameter file identifiable for a given run
#Can leave it blank if you don't care.
paramfiletag = 'testgrid'

#No need to add an underscore/jobnumber, the script will do that for you.
labelend = 'test'

#What number to start counting from, must be an integer
jobnumstart = 1

#Define parameters to feed into file, must be filled with at least 1 value
#Want to check that the values for amaxs and epsilon are possible in the sample job file

mstar   = [1.3] #Mass of the star in Msun
tstar   = [4730] #Temperature of the star in K
rstar   = [1.6] #Radius of the star in solar radii
dist    = [140] #Distance to the star in pc
mdot    = [3.3e-9] #Mass accretion rate in the disk in solar masses per year

amaxs   = [0.25] #Maximum grain size in the upper layers of the disk NOTE: Only acceptss certain values, see the docstring above
epsilon = [0.1] #Settling parameter
ztran   = [0.1] #height of transition between big and small grains, in hydrostatic scale heights
alpha   = [1e-2] # Viscosity in the disk
rdisk   = [200] #Outer radius of the disk

rc      = [2000.] #Critial radius of tapered edge. If 2000, it will not have any effect.
gamma   = [2.0] #Exponent of tapered edge.

temp    = [1400] #Sublimation temperature of the grains
altinh  = [1] #Height of the disk in scale heights. Often better to leave as 1 and scale it later

lamaxb  = ['1mm'] #Maximum grain size in the midplane NOTE: Only accepts certain values
mui     = [0.5] #Cosine of the inclination of the disk.
tshock  = [8000] #Temperature of the shock. Usually left at 8000K

fracolive = [1.0] #Fraction of olivines
fracpyrox = [0.0] #Fraction of pyroxine
fracforst = [0.0] #Fraction of fosterite
fracent   = [0.0] #Fraction of enstatite

d2g = [0.0065] #Dust to gas mass ratio.

#Gap creator:
imod = False # If True, the disk structure will be modified as chosen with the following parameters.
#NOTE: these are python lists, so should be written to the jobfile as "[XX]"
#All used deltas need to have the same length as inter_r, so be careful with the combinations.
inter_r = ['[]']
rho_deltas = ['[]']
temp_deltas = ['[]']
epsbig_deltas = ['[]']
eps_deltas = ['[]']

#OPTIONAL PARAMETERS:
#If you want them to be the same as their associated parameters (amaxs, mdot) then set them to [None]
#Don't forget the brackets.
mdotstar = [None] #Mass accretion rate at the star
amaxw   = [None] #Maximum grain size at the wall

fill = 3 #Zero padding for the job numbers. Unless you are running many models (>1000) 3 is standard.

iwall = False #Set this to True if you want ONLY walls. (i.e., turns off the disk code).

#***********************************************
#Unlikely you need to change anything below here.
#***********************************************

#Open up a file and print the parameter names
f = open(gridpath+paramfiletag+'job_params.txt', 'w')
f.writelines('Job Number, amaxs, amaxw, epsilon, ztran, mstar, tstar, rstar, \
dist, mdot, mdotstar, tshock, alpha, mui, rdisk, rc, gamma, temp, altinh, \
fracolive, fracpyrox, fracforst, fracent, lamaxb, d2g, inter_r, rho_deltas, \
temp_deltas, epsbig_deltas, eps_deltas \n')

#Write each iteration as a row in the table
for ind, values in enumerate(itertools.product(amaxs, amaxw, epsilon, ztran,
mstar, tstar, rstar, dist, mdot, mdotstar, tshock, alpha, mui, rdisk, rc,
gamma, temp, altinh, fracolive, fracpyrox, fracforst, fracent, lamaxb, d2g,
inter_r, rho_deltas, temp_deltas, epsbig_deltas, eps_deltas)):

    #Handle the cases of mdotstar and amaxw which are optional parameters
    #NOTE: DO NOT CHANGE THE PARAMETERS BEFORE MDOTSTAR OR THIS WILL BREAK! YOU HAVE BEEN WARNED
    values = np.array(values)
    if mdotstar == [None]:
        values[9] = values[8]
    if amaxw == [None]:
        values[1] = values[0]
    values = tuple(values)

    f.writelines(str(ind+jobnumstart)+', '+str(values)[1:-1]+ '\n')
f.close()

#Open up the table
table = ascii.read(gridpath+paramfiletag+'job_params.txt')

#Throw up a warning if you are making a huge grid
if len(table) > 1000:
    print('WARNING! GRID IS LARGE AND WILL USE SIGNIFICANT LIMITED COMPUTING RESOURCES! SPEAK TO A SUPERVISOR BEFORE RUNNING ON THE CLUSTER!!!')
    warning = input('Please enter "I understand" to continue making this grid: ')

    if warning != 'I understand':
        raise ValueError('Job creation canceled.')
    else:
        print('Continuing job creation process...')

#Create the jobfiles using edge.job_file_create
for i in range(len(table)):
    label = labelend+'_'+str(i+jobnumstart).zfill(fill)

    # Remove the apostrophes
    if "'" in table['lamaxb'][i]:
        table['lamaxb'][i] = table['lamaxb'][i].split("'")[1]
    if "'" in table['inter_r'][i]:
        table['inter_r'][i] = table['inter_r'][i].split("'")[1]
    if "'" in table['rho_deltas'][i]:
        table['rho_deltas'][i] = table['rho_deltas'][i].split("'")[1]
    if "'" in table['temp_deltas'][i]:
        table['temp_deltas'][i] = table['temp_deltas'][i].split("'")[1]
    if "'" in table['epsbig_deltas'][i]:
        table['epsbig_deltas'][i] = table['epsbig_deltas'][i].split("'")[1]
    if "'" in table['eps_deltas'][i]:
        table['eps_deltas'][i] = table['eps_deltas'][i].split("'")[1]

    edge.job_file_create(i+jobnumstart, gridpath, \
    amaxs     = table['amaxs'][i],\
    eps       = table['epsilon'][i],\
    ztran     = table['ztran'][i],\
    mstar     = table['mstar'][i], \
    tstar     = table['tstar'][i], \
    rstar     = table['rstar'][i], \
    dist      = table['dist'][i], \
    mdot      = table['mdot'][i], \
    mdotstar  = table['mdotstar'][i], \
    tshock    = table['tshock'][i], \
    alpha     = table['alpha'][i], \
    mui       = table['mui'][i], \
    rdisk     = table['rdisk'][i], \
    rc        = table['rc'][i], \
    gamma     = table['gamma'][i], \
    temp      = table['temp'][i], \
    altinh    = table['altinh'][i],\
    fracolive = table['fracolive'][i], \
    fracpyrox = table['fracpyrox'][i], \
    fracforst = table['fracforst'][i], \
    fracent   = table['fracent'][i], \
    amaxw     = table['amaxw'][i],\
    lamaxb    = table['lamaxb'][i],\
    d2g       = table['d2g'][i],\
    inter_r   = table['inter_r'][i],\
    rho_deltas= table['rho_deltas'][i],\
    temp_deltas= table['temp_deltas'][i],\
    epsbig_deltas= table['epsbig_deltas'][i],\
    eps_deltas= table['eps_deltas'][i],\
    fill      = fill,\
    iwall     = iwall,\
    imod      = imod,\
    labelend  = label)

#Make a run all file for the jobs you just created
edge.create_runall(jobnumstart, jobnumstart+len(table)-1, clusterpath, optthin = False, outpath = gridpath, fill = fill)

max_stars_repo_name: zihuaxin/EDGE

max_issues_repo_name: zihuaxin/EDGE

max_forks_repo_name: zihuaxin/EDGE

reason: import numpy,from astropy


text:
from abess import *
from utilities import *
import numpy as np

class TestCheck:
    """
    Test for argument error, which should be recognized before the algorithm.
    """

    @staticmethod
    def test_fit():
        # path
        try:
            model = abessLm(path_type='other')
            model.fit([[1]], [1])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model = abessLm(support_size=[3])
            model.fit([[1]], [1])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model = abessLm(path_type='gs', s_min=1, s_max=0)
            model.fit([[1]], [1])
        except ValueError as e:
            print(e)
        else:
            assert False

        # ic
        try:
            model = abessLm(ic_type='other')
            model.fit([[1]], [1])
        except ValueError as e:
            print(e)
        else:
            assert False

        # exchange_num
        try:
            model = abessLm(exchange_num=-1)
            model.fit([[1]], [1])
        except ValueError as e:
            print(e)
        else:
            assert False

        # screening_size
        try:
            model = abessLm(screening_size=3)
            model.fit([[1]], [1])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model = abessLm(support_size=[2],
                            screening_size=1)
            model.fit([[1, 2, 3]], [1])
        except ValueError as e:
            print(e)
        else:
            assert False

        # primary_fit_xxx
        try:
            model = abessLogistic(primary_model_fit_max_iter=0.5)
            model.fit([[1]], [1])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model = abessLogistic(primary_model_fit_epsilon=-1)
            model.fit([[1]], [1])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model = abessLogistic(primary_model_fit_epsilon=-1)
            model.fit([[1]], [1])
        except ValueError as e:
            print(e)
        else:
            assert False

        # thread
        try:
            model = abessLm(thread=-1)
            model.fit([[1]], [1])
        except ValueError as e:
            print(e)
        else:
            assert False

        # splicing_type
        try:
            model = abessLm(splicing_type=-1)
            model.fit([[1]], [1])
        except ValueError as e:
            print(e)
        else:
            assert False

        # cv
        try:
            model = abessLm(cv=2)
            model.fit([[1]], [1])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model = abessLm(cv=2)
            cv_fold_id = [1, 1]
            model.fit([[1], [1]], [1, 1], cv_fold_id=cv_fold_id)
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model = abessLm(cv=2)
            cv_fold_id = [1, 2, 1]
            model.fit([[1], [1]], [1, 1], cv_fold_id=cv_fold_id)
        except ValueError as e:
            print(e)
        else:
            assert False

        model = abessLm()
        # datatype error
        try:
            model.fit([['c', 1, 1]], [1])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model.fit([[1, 1, 1]], [1], weight=['c'])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model1 = abessLm(cv='c')
            model1.fit([[1]], [1])
        except ValueError as e:
            print(e)
        else:
            assert False

        # incompatible shape
        try:
            model.fit([1, 1, 1], [1])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model.fit([[1, 1, 1]], [1, 2])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model.fit([[1, 1, 1]], [1])
            model.predict([[1, 1]])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model.fit([[1, 1, 1]], [1], weight=[1, 2])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model.fit([[1, 1, 1]], [1], group=[1])
        except ValueError as e:
            print(e)
        else:
            assert False

        # lack of necessary parameter
        try:
            model.fit(X=[[1]])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model.fit(y=[1])
        except ValueError as e:
            print(e)
        else:
            assert False

    @staticmethod
    def test_pca():
        """
        For `abess.pca.abessPCA`.
        """
        model = abessPCA()
        # datatype error
        try:
            model.fit([['c']])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model.fit(Sigma=[['c']])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model.fit(Sigma=[[np.nan]])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model1 = abessPCA(cv='c')
            model1.fit([[1]])
        except ValueError as e:
            print(e)
        else:
            assert False

        # incompatible shape
        try:
            model.fit([1])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model.fit(Sigma=[1])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model.fit([[1]], group=[1, 2])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model1 = abessPCA(support_size=np.array([1, 2]))
            model1.fit([[1]])
        except ValueError as e:
            print(e)
        else:
            assert False

        # lack of necessary parameter
        try:
            model.fit()
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model1 = abessPCA(cv=5)
            model1.fit(Sigma=[[1]])
        except ValueError as e:
            print(e)
        else:
            assert False

        # number
        try:
            model.fit([[1]], [1], number=-1)
        except ValueError as e:
            print(e)
        else:
            assert False

        # invalid sigma
        try:
            model.fit(Sigma=[[1, 0], [1, 0]])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model.fit(Sigma=[[-1, 0], [0, -1]])
        except ValueError as e:
            print(e)
        else:
            assert False

        # invalid arg
        try:
            model1 = abessPCA(ic_type='other')
            model1.fit([[1]])
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model1 = abessPCA(cv=5)
            model1.fit([[1]])
        except ValueError as e:
            print(e)
        else:
            assert False

    @staticmethod
    def test_rpca():
        model = abessRPCA()
        # datatype error
        try:
            model.fit([['c']], r=1)
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model.fit([[1]], r='c')
        except ValueError as e:
            print(e)
        else:
            assert False

        # incompatible shape
        try:
            model.fit([1], r=1)
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model.fit(1, r=1)
        except ValueError as e:
            print(e)
        else:
            assert False

        try:
            model.fit([[1]], r=1, group=[1, 2])
        except ValueError as e:
            print(e)
        else:
            assert False

        # invalid arg
        try:
            model1 = abessRPCA(ic_type='other')
            model1.fit([[1]], r=1)
        except ValueError as e:
            print(e)
        else:
            assert False


text:
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests and benchmarks for interacting with the `tf.compat.v1.Session`."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import time

import numpy as np

from six.moves import xrange  # pylint: disable=redefined-builtin
from tensorflow.python.client import session
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.training import server_lib

class SessionBenchmark(test.Benchmark):
  """Tests and benchmarks for interacting with the `tf.compat.v1.Session`."""

  def _benchmarkFeed(self, name, target, size, iters):
    """Runs a microbenchmark to measure the cost of feeding a tensor.

    Reports the median cost of feeding a tensor of `size` * `sizeof(float)`
    bytes.

    Args:
      name: A human-readable name for logging the output.
      target: The session target to use for the benchmark.
      size: The number of floating-point numbers to be feed.
      iters: The number of iterations to perform.
    """
    feed_val = np.random.rand(size).astype(np.float32)
    times = []
    with ops.Graph().as_default():
      p = array_ops.placeholder(dtypes.float32, shape=[size])
      # Fetch the operation rather than the tensor, to avoid measuring the time
      # to fetch back the value.
      no_op = array_ops.identity(p).op
      with session.Session(target) as sess:
        sess.run(no_op, feed_dict={p: feed_val})  # Warm-up run.
        for _ in xrange(iters):
          start_time = time.time()
          sess.run(no_op, feed_dict={p: feed_val})
          end_time = time.time()
          times.append(end_time - start_time)
    print("%s %d %f" % (name, size, np.median(times)))
    self.report_benchmark(iters=1, wall_time=np.median(times), name=name)

  def _benchmarkFetch(self, name, target, size, iters):
    """Runs a microbenchmark to measure the cost of fetching a tensor.

    Reports the median cost of fetching a tensor of `size` * `sizeof(float)`
    bytes.

    Args:
      name: A human-readable name for logging the output.
      target: The session target to use for the benchmark.
      size: The number of floating-point numbers to be fetched.
      iters: The number of iterations to perform.
    """
    times = []
    with ops.Graph().as_default():
      # Define the tensor to be fetched as a variable, to avoid
      # constant-folding.
      v = variables.Variable(random_ops.random_normal([size]))
      with session.Session(target) as sess:
        sess.run(v.initializer)
        sess.run(v)  # Warm-up run.
        for _ in xrange(iters):
          start_time = time.time()
          sess.run(v)
          end_time = time.time()
          times.append(end_time - start_time)
    print("%s %d %f" % (name, size, np.median(times)))
    self.report_benchmark(iters=1, wall_time=np.median(times), name=name)

  def _benchmarkFetchPrebuilt(self, name, target, size, iters):
    """Runs a microbenchmark to measure the cost of fetching a tensor.

    Reports the median cost of fetching a tensor of `size` * `sizeof(float)`
    bytes.

    Args:
      name: A human-readable name for logging the output.
      target: The session target to use for the benchmark.
      size: The number of floating-point numbers to be fetched.
      iters: The number of iterations to perform.
    """
    times = []
    with ops.Graph().as_default():
      # Define the tensor to be fetched as a variable, to avoid
      # constant-folding.
      v = variables.Variable(random_ops.random_normal([size]))
      with session.Session(target) as sess:
        sess.run(v.initializer)
        runner = sess.make_callable(v)
        runner()  # Warm-up run.
        for _ in xrange(iters):
          start_time = time.time()
          runner()
          end_time = time.time()
          times.append(end_time - start_time)
    print("%s %d %f" % (name, size, np.median(times)))
    self.report_benchmark(iters=1, wall_time=np.median(times), name=name)

  def _benchmarkRunOp(self, name, target, iters):
    """Runs a microbenchmark to measure the cost of running an op.

    Reports the median cost of running a trivial (Variable) op.

    Args:
      name: A human-readable name for logging the output.
      target: The session target to use for the benchmark.
      iters: The number of iterations to perform.
    """
    times = []
    with ops.Graph().as_default():
      # Define the op to be run as a variable, to avoid
      # constant-folding.
      v = variables.Variable(random_ops.random_normal([]))
      with session.Session(target) as sess:
        sess.run(v.initializer)
        sess.run(v.op)  # Warm-up run.
        for _ in xrange(iters):
          start_time = time.time()
          sess.run(v.op)
          end_time = time.time()
          times.append(end_time - start_time)
    print("%s %f" % (name, np.median(times)))
    self.report_benchmark(iters=1, wall_time=np.median(times), name=name)

  def _benchmarkRunOpPrebuilt(self, name, target, iters):
    """Runs a microbenchmark to measure the cost of running an op.

    Reports the median cost of running a trivial (Variable) op.

    Args:
      name: A human-readable name for logging the output.
      target: The session target to use for the benchmark.
      iters: The number of iterations to perform.
    """
    times = []
    with ops.Graph().as_default():
      # Define the op to be run as a variable, to avoid
      # constant-folding.
      v = variables.Variable(random_ops.random_normal([]))
      with session.Session(target) as sess:
        sess.run(v.initializer)
        runner = sess.make_callable(v.op)
        runner()  # Warm-up run.
        for _ in xrange(iters):
          start_time = time.time()
          runner()
          end_time = time.time()
          times.append(end_time - start_time)
    print("%s %f" % (name, np.median(times)))
    self.report_benchmark(iters=1, wall_time=np.median(times), name=name)

  def benchmarkGrpcSession(self):
    server = server_lib.Server.create_local_server()
    self._benchmarkFeed("benchmark_session_feed_grpc_4B", server.target, 1,
                        30000)
    session.Session.reset(server.target)
    self._benchmarkFeed("benchmark_session_feed_grpc_4MB", server.target,
                        1 << 20, 25000)
    session.Session.reset(server.target)
    self._benchmarkFetch("benchmark_session_fetch_grpc_4B", server.target, 1,
                         40000)
    session.Session.reset(server.target)
    self._benchmarkFetch("benchmark_session_fetch_grpc_4MB", server.target,
                         1 << 20, 20000)
    session.Session.reset(server.target)
    self._benchmarkFetchPrebuilt("benchmark_session_fetchprebuilt_grpc_4B",
                                 server.target, 1, 50000)
    session.Session.reset(server.target)
    self._benchmarkFetchPrebuilt("benchmark_session_fetchprebuilt_grpc_4MB",
                                 server.target, 1 << 20, 50000)
    session.Session.reset(server.target)
    self._benchmarkRunOp("benchmark_session_runop_grpc", server.target, 50000)
    session.Session.reset(server.target)
    self._benchmarkRunOpPrebuilt("benchmark_session_runopprebuilt_grpc",
                                 server.target, 100000)
    session.Session.reset(server.target)

  def benchmarkDirectSession(self):
    self._benchmarkFeed("benchmark_session_feed_direct_4B", "", 1, 80000)
    self._benchmarkFeed("benchmark_session_feed_direct_4MB", "", 1 << 20, 20000)
    self._benchmarkFetch("benchmark_session_fetch_direct_4B", "", 1, 100000)
    self._benchmarkFetch("benchmark_session_fetch_direct_4MB", "", 1 << 20,
                         20000)
    self._benchmarkFetchPrebuilt("benchmark_session_fetchprebuilt_direct_4B",
                                 "", 1, 200000)
    self._benchmarkFetchPrebuilt("benchmark_session_fetchprebuilt_direct_4MB",
                                 "", 1 << 20, 200000)
    self._benchmarkRunOp("benchmark_session_runop_direct", "", 200000)
    self._benchmarkRunOpPrebuilt("benchmark_session_runopprebuilt_direct", "",
                                 200000)

if __name__ == "__main__":
  test.main()

max_stars_repo_name: uve/tensorflow

max_issues_repo_name: uve/tensorflow

max_forks_repo_name: uve/tensorflow


text:
import sys
import string
import warnings
import math
import os
import random
from collections import Counter

from multiprocessing import Pool

import torch
import torchtext

import numpy as np
import pandas as pd

import av

import lxml, lxml.html
import html

import spacy
import youtube
import pytube

from tqdm import tqdm

class YouTubeDataset(torch.utils.data.IterableDataset):
    """
    YouTubeDataset is a flexible video dataset creation tool.
    It allows the downloading of a complete YouTube channel as a dataset.

    Downloading a channel dataset with YouTubeDataset is easy, set the ``download`` flag to True,
    and set ``api_key`` to a Google Cloud Platform API Key with rights for the YouTube Data API V3,
    and the YouTube ``channel_id``. The download will start automatically, and cache the dataset
    locally in ``root_dir``/``channel``.

    To use the dataset, set the `fields` parameter to a list of fields that will be extracted
    from the dataset and presented by the iterator, and the ``key`` parameter to what key field is.

    ``fields`` is a list of one or more of the following ``fieldnames``:

    ``key`` refers to the dataset key field. The key field determines the length of all the
    other dataset fields.
    If ``key`` is ``VIDEO``, then the duration of ``AUDIO`` and ``TEXT`` fields are
    determined by the ``video_cliplen`` which is the number of frames to present.

    If using ``CC_TEXT`` as the key, you will need to get the max text duration from the ``max_text_dur`` property,
    and pad the data to the max duration in the video_transform and audio_transform.

    Parameters:
        root_dir (string): Root directory of the YouTube Dataset.
        channel (string):  Name to identify dataset locally
        split (string): The split to present - splits are determined by ``splits`` dictionary
        fields (list, optional): A list of ``fieldnames`` representing the data fields presented by the iterator
        key (string, optional): The key field's ``fieldname``
        video_format (string, optional): The video pixel format one of: 'yuv420p','rgb24', 'rgba'
        video_framerate (int, optional): if specified, it will resample the video
            so that it has `frame_rate`, and then the video_cliplen will be defined on the resampled video
        video_cliplen (int, optional): number of frames in a clip
        video_stridelen (int, tuple, optional): number of frames between each clip, if tuple a range for the stride (1,199) random between 1 and 199
        video_transform (callable, optional): A function/transform that  takes in a TxHxWxC video
            and returns a transformed version.
        image_transform (callable, optional): A function/transform that  takes in a CxHxW image
            and returns a transformed version.
        audio_format (string): The audio sample format, one of: 'u8', 's16', 's32', 'flt','dbl'
        audio_layout (string, int, optional): The audio channel layout, either an integer number of channels, or
            audio_layout can be one or several of the following notations,
            separated by '+' or '|':
            - the name of an usual channel layout (mono, stereo, 4.0, quad, 5.0,
              5.0(side), 5.1, 5.1(side), 7.1, 7.1(wide), downmix);
            - the name of a single channel (FL, FR, FC, LFE, BL, BR, FLC, FRC, BC,
              SL, SR, TC, TFL, TFC, TFR, TBL, TBC, TBR, DL, DR);
            - a number of channels, in decimal, followed by 'c', yielding
              the default channel layout for that number of channels (@see
              av_get_default_channel_layout);
            - a channel layout mask, in hexadecimal starting with "0x" (see the
              AV_CH_* macros).
            Example: "stereo+FC" = "2c+FC" = "2c+1c" = "0x7"
        audio_sr (int, optional): The sample rate of audio clips in samples per second.
        audio_cliplen (int, optional): number of samples in a clip
        audio_transform (callable, optional): A function/transform that takes in a KxL audio tensor
            and returns a transformed version.
        text_lang (string, optional): the language of the captions using ISO 639-1. Default "en".
        text_transform (callable, optional): A function/transform that takes in a string of text
            and returns the embedded text as a Tensor of any size.
        dataset_fraction (float): fractuion of dataset to use (0. - 1.)
    Returns:
        tuple of data fields from dataset filtered by ``fields`` and ``key``.
        ``fields`` list items will return data in the following formats:
            ``VIDEOID`` (string): The YouTube video ID
            ``TITLE`` (string): The title of the video
            ``DESCRIPTION`` (string): The description of the video
            ``TIME`` (float): The presentation time in seconds of the returned data ``key``
            ``DURATION`` (float): The duration of the returned data ``key`` in seconds
            ``CC_TEXT`` (string): Close Caption Text
            ``IF_DATA`` (Tensor[C, H, W]): Image Tensor 'H' - Height, 'W' - Width, 'C' Channels
            ``VF_DATA`` (Tensor[T, H, W, C]): 'T' Video Frames or single frame
            ``AF_DATA`` (Tensor[K, L]): Audio Frames where `K` is the number of audio channels, and 'L' is samples
    """

    F_VIDEOID = 'VIDEOID'
    F_TITLE = 'TITLE'
    F_DESCRIPTION = 'DESCRIPTION'
    F_TIME = 'TIME'
    F_DURATION = 'DURATION'
    F_VF_DATA = 'VIDEO'
    F_IF_DATA = 'IMAGE'
    F_AF_DATA = 'AUDIO'
    F_CC_TEXT = 'TEXT'

    def __init__(self, root_dir, channel, split='train', fields=['IMAGE','TIME'], key='IMAGE',
                 video_format='rgb24', video_framerate=None, video_cliplen=None, video_stridelen=None, image_transform=None, video_transform=None,
                 audio_format='flt', audio_layout='mono', audio_sr=16000,  audio_cliplen=None, audio_transform=None,
                 text_lang='en', text_transform=None, dataset_fraction=1.0,
                 download=False, api_key=None, channel_id=None, user_name=None, splits={'train':0.80, 'test':0.20}):

        assert (key in fields), "YouTubeDataset: fields must contain key"
        assert not download or (download and api_key != None), "YouTubeDataset: download requires api_key from Google API Console with Youtube access"
        assert not download or (download and (channel_id != None or user_name != None)), "YouTubeDataset: download requires one of channel_id or user_name"
        assert key in fields, "YouTubeDataset: key must be in fields"
        assert (key == YouTubeDataset.F_VF_DATA and video_cliplen > 1) or True, "YouTubeDataset: video_cliplen must be 1 or greater if key is 'VIDEO'"
        assert (key == YouTubeDataset.F_AF_DATA and audio_cliplen > 1) or True, "YouTubeDataset: audio_cliplen must be 1 or greater if key is 'AUDIO'"

        if download:
            dl = YouTubeDatasetDownloader(root_dir, channel, splits, api_key, caption_lang=text_lang, channel_id=channel_id, user_name=user_name)
            dl.download_sync()

        csv_file = os.path.join(root_dir, channel, split + '.csv')

        assert os.path.isfile(csv_file), "YouTubeDataset: File Not Found " + csv_file

        self.dataframe = pd.read_csv(csv_file)
        self.root_dir = root_dir
        self.split = split
        self.fields = fields
        self.key = key
        self.channel = channel
        self.image_transform = image_transform
        self.video_transform = video_transform
        self.audio_transform = audio_transform
        self.text_transform = text_transform
        self.text_lang = text_lang

        if dataset_fraction > 0.0 and dataset_fraction < 1.0:
            self.dataframe=self.dataframe.sample(frac=dataset_fraction)

        self.start = 0
        self.size = self.dataframe.shape[0]
        self.video_format = video_format
        self.video_framerate = video_framerate
        self.video_cliplen = video_cliplen
        self.video_stridelen=video_stridelen
        self.audio_format = audio_format
        self.audio_sr = audio_sr
        self.audio_layout = audio_layout
        self.audio_cliplen = audio_cliplen

        self.raw_text_cache = None

        self._len = None

    def __iter__(self):
        return YTDSIterator(self)

    @property
    def length(self):
        if self._len == None:
            if self.key == 'TEXT':
                self._len = len(self.raw_text)
            elif self.key == 'VIDEO' or self.key == 'IMAGE':
                if self.video_stridelen is None:
                    stride = self.video_cliplen
                elif isinstance(self.video_stridelen, tuple):
                    stride = random.randint(*self.video_stridelen)
                else:
                    stride = self.video_stridelen

                self._len = int((sum([dur for pts, dur, text in self.raw_text]) * (self.video_framerate if self.video_framerate != None else 24)) // stride)
            elif self.key == 'AUDIO':
                self._len = int((sum([dur for pts, dur, text in self.raw_text]) * (self.video_framerate if self.video_framerate != None else 24)) //
                                (self.audio_cliplen if self.audio_cliplen != None else 512))
        return self._len

    @property
    def fieldnames(self):
        """Get the list of valid ``fieldnames``.
        Returns:
            list: list of valid fieldnames
        """
        return [YouTubeDataset.F_VIDEOID, YouTubeDataset.F_TITLE, YouTubeDataset.F_DESCRIPTION, YouTubeDataset.F_TIME,
                YouTubeDataset.F_DURATION, YouTubeDataset.F_VF_DATA, YouTubeDataset.F_IF_DATA, YouTubeDataset.F_AF_DATA,
                YouTubeDataset.F_CC_TEXT]

    @property
    def raw_text(self):
        """Returns the raw TEXT of the dataset as a list
        Returns:
            list: [(TIME, DURATION, TEXT) ...] the Caption portion of the dataset
        """
        if self.raw_text_cache == None:
            self.raw_text_cache = [ (pts, dur, text) for pts, dur, text in YouTubeDataset(self.root_dir, self.channel, self.split,
                                                             fields=[YouTubeDataset.F_TIME,YouTubeDataset.F_DURATION,YouTubeDataset.F_CC_TEXT],
                                                             key=YouTubeDataset.F_CC_TEXT) ]
        return self.raw_text_cache

    @property
    def max_text_dur(self):
        """Returns the max duration of the TEXT fields.
        Returns:
            float: duration in seconds
        """

        return max([dur for pts, dur, text in self.raw_text])

    @property
    def max_text_len(self):
        """Returns the max length of the TEXT fields in chars.
        Returns:
            int: length in characters
        """
        return max([len(text) for pts, dur, text in self.raw_text])

    def max_text_tokens(self,  tokenizer, vocab=None):
        """Returns the max number of tokens in the TEXT fields, using the specified totchtext tokenizer
        Returns:
            int: length in tokens
        """
        if vocab is None:
            return max([len(tokenizer(text)) for pts, dur, text in self.raw_text])
        else:
            return max([len([t for t in tokenizer(text) if str(t) in vocab.stoi]) for pts, dur, text in self.raw_text])

    @staticmethod
    def build_vocab(root_dir, channel, vocab='vocab.txt', language="en", pretrained="fasttext.en.300d", pos=None, lemmatize=False):
        """Returns a new torchtext.experimental.vectors.Vectors object for the dataset, either from the vectors file or
           building a new one from the dataset using the tokenizer and the pretrained vectors
        Params:
            vocab: file name in the dataset directory of the cached vocab
            tokenizer: tokenizer function
            pretrained: one of available pretrained vectors: charngram.100d fasttext.en.300d fasttext.simple.300d glove.42B.300d
                        glove.840B.300d glove.twitter.27B.25d glove.twitter.27B.50d glove.twitter.27B.100d glove.twitter.27B.200d g
                        love.6B.50d glove.6B.100d glove.6B.200d glove.6B.300d
            pos: parts of speech to include in vocab, None for all words or, list of Universal POS tags:
                    'ADJ','ADP','PUNCT','ADV','AUX','SYM','INTJ','CONJ','X','NOUN','DET','PROPN','NUM','VERB','PART','PRON','SCONJ'
        Returns:
            torchtext.experimental.vectors.Vectors
        """
        cache_path = os.path.join(root_dir, channel, vocab)

        if not os.path.exists(cache_path):
            assert pretrained in torchtext.vocab.pretrained_aliases, "build_vector: allowed pretrained vectors values are " + str(torchtext.vocab.pretrained_aliases.keys())
            counter = Counter()
            pt = torchtext.vocab.pretrained_aliases[pretrained]()

            nlp = spacy.load(language)

            if pos is None and not lemmatize:
                tokenizer = lambda text: [t.text for t in nlp.tokenizer(text) if t.text in pt.stoi]
            else:
                tokenizer = lambda text: [t.text if not lemmatize else t.lemma_ for t in nlp(text) if t.text in pt.stoi and t.pos_ in pos]

            for split in ['train', 'test', 'validate']:
                for [text] in tqdm(YouTubeDataset(root_dir, channel, split, fields=[YouTubeDataset.F_CC_TEXT], key=YouTubeDataset.F_CC_TEXT), desc=split):
                    counter.update(tokenizer(text))

            toks = {tok: ix for ix, tok in enumerate(counter)}
            vecs = [pt[tok] for tok in counter]

            with open(cache_path, 'w') as fp:
                fp.writelines([",".join([t, str(counter[t]), " ".join([*[str(float(d)) for d in vecs[toks[t]]], "\n"])]) for t in toks])

        # Open the Vector
        with open(cache_path, 'r') as fp:
            toks={}
            vecs=[]
            counter=Counter()

            for ix, line in enumerate(fp.readlines()):
                key,count,vec = line.split(',')
                toks[key] = ix
                counter[key] = int(count)
                vecs.append(torch.tensor([float(v) for v in vec.split(' ') if v != "\n"]))
            vocab = torchtext.vocab.Vocab(counter)
            vocab.set_vectors(toks, vecs, 300)
            return vocab

    @staticmethod
    def worker_init(worker_id):
        """Worker Init Function for DataLoader. See torch.utils.data.DataLoader
        Parameters:
            worker_id (int): The ID of the worker
        """
        worker_info = torch.utils.data.get_worker_info()
        dataset = worker_info.dataset  # the dataset copy in this worker process
        overall_start = dataset.start
        overall_size = dataset.size
        worker_id = worker_info.id

        if overall_size < worker_info.num_workers:
            raise ValueError('Number of workers must be less than number of videos in split ('+str(overall_size)+')')

        # configure the dataset to only process the split workload
        per_worker = overall_size // worker_info.num_workers
        dataset.start = overall_start + worker_id * per_worker
        dataset.size = min(per_worker, overall_size - dataset.start)

class CCFrame():
    def __init__(self, time, duration, text):
        self.__time = time
        self.__duration = duration
        self.__text = text

    @property
    def time(self):
        return self.__time

    @property
    def duration(self):
        return self.__duration

    @property
    def text(self):
        return self.__text

    def __repr__(self):
        return '<%s pts=%d, duration=%d, text=%s at 0x%x>' % (
            self.__class__.__name__,
            self.time,
            self.duration,
            self.text,
            id(self)
        )

    @staticmethod
    def ParseCC(caption_path):
        with open(caption_path, 'r') as myfile:
            data = myfile.read()

        forest = lxml.html.fragments_fromstring(data, no_leading_text=True)

        transcript = []

        for e in forest:
            if e.tag == 'transcript':
                for text in e:
                    if text.tag == 'text':
                        try:
                            transcript.append(CCFrame(float(text.attrib['start']), float(text.attrib['dur']), html.unescape(text.text)))
                        except:
                            pass

        return transcript

# This dataset returns the following tuple for each iteration
# pts as float
# video_frame as numpy array - the video frame at this pts
# audio_frame as numpy array - the entire audio frame for this pts
# caption_frame as string    - the caption text for this pts

class YTDSIterator():
    def __init__(self, dataset):
        self.dataset = dataset

        self.active_ix = -1
        self.active = None

        self.next_active()

    def next_active(self):
        try:
            # Get The Active Record from the dataset
            if self.active_ix == -1:
                self.active_ix = self.dataset.start
            else:
                self.active_ix = self.active_ix + 1

            if self.active_ix >= (self.dataset.start + self.dataset.size):
                raise StopIteration()

            self.active = self.dataset.dataframe.iloc[self.active_ix]

            self.video_id = self.active['video_id']
            self.title = self.active['title']
            self.description = self.active['description']

            video_path = os.path.join(self.dataset.root_dir,self.dataset.channel, 'video', self.video_id + '.mp4')
            caption_path = os.path.join(self.dataset.root_dir,self.dataset.channel, 'caption', self.video_id + '.xml')

            # Initialize Video
            if (YouTubeDataset.F_VF_DATA in self.dataset.fields or
                YouTubeDataset.F_IF_DATA in self.dataset.fields):
                self.vcontainer = av.open(video_path)
                self.video = self.vcontainer.decode(video=0)
                self.video_rate = self.vcontainer.streams.video[0].average_rate
            else:
                self.vcontainer = None
                self.video = None
                self.video_rate = 0

            # Initialize Audio
            if (YouTubeDataset.F_AF_DATA in self.dataset.fields):
                self.acontainer = av.open(video_path)
                self.audio = self.acontainer.decode(audio=0)
                if (self.acontainer.streams.audio[0].sample_rate != self.dataset.audio_sr or
                    self.acontainer.streams.audio[0].layout != self.dataset.audio_layout or
                    self.acontainer.streams.audio[0].format != self.dataset.audio_format):
                    self.audio_resampler = av.audio.resampler.AudioResampler(self.dataset.audio_format,
                                                                             self.dataset.audio_layout, self.dataset.audio_sr)
                else:
                    self.audio_resampler = None
            else:
                self.acontainer = None
                self.audio = None
                self.audio_resampler = None

            # init captions
            if YouTubeDataset.F_CC_TEXT in self.dataset.fields:
                self.caption = iter(CCFrame.ParseCC(caption_path))
            else:
                self.caption = None

            # initialize VideoFrame start
            self.vclip = []
            self.vclip_time = 0.0

            # Initialize Audioframe Start
            self.aclip_time = 0.0
            self.aclip = []

            # Initialize Caption Start
            self.cclip_time = 0.0
            self.cclip_duration = 0.0
            self.cclip = []

            # set up the result
            self.result = {
                YouTubeDataset.F_VIDEOID: self.video_id,
                YouTubeDataset.F_TITLE: self.title,
                YouTubeDataset.F_DESCRIPTION: self.description
            }
        except StopIteration as si:
            raise si

        except Exception as e:
            print('Exception', type(e), e.args, e)
            raise StopIteration()

    def __iter__(self):
        return self

    def __next__(self):
        try:
            result = self.result

            if self.dataset.key == YouTubeDataset.F_CC_TEXT:
                clip_duration = self.cclip_duration
                clip_time = self.cclip_time

            if self.dataset.key == YouTubeDataset.F_VF_DATA:
                if self.dataset.video_framerate == None:
                    clip_duration = self.dataset.video_cliplen / self.video_rate
                else:
                    clip_duration = self.dataset.video_cliplen / self.dataset.video_framerate
                clip_time = self.vclip_time

            if self.dataset.key == YouTubeDataset.F_IF_DATA:
                clip_duration = 1. / self.video_rate
                clip_time = self.vclip_time

            if self.dataset.key == YouTubeDataset.F_AF_DATA:
                clip_duration = self.dataset.audio_cliplen / self.dataset.audio_sr
                clip_time = self.aclip_time

            if YouTubeDataset.F_CC_TEXT in self.dataset.fields:

                if self.dataset.key == YouTubeDataset.F_CC_TEXT:
                    # Get Next Caption
                    cc = next(self.caption)
                    self.cclip = [cc]
                    self.cclip_time = clip_time = cc.time
                    self.cclip_duration = clip_duration = cc.duration
                    cc_text = cc.text

                    result[YouTubeDataset.F_TIME] = self.cclip_time
                    result[YouTubeDataset.F_DURATION] = self.cclip_duration

                else:
                    # remove old text clips
                    for c in self.cclip:
                        if (c.time + c.duration) < clip_time:
                            self.cclip.remove(c)
                            if len(self.cclip) > 0:
                                self.cclip_time = self.cclip[0].time
                            else:
                                self.cclip_time = 0
                            self.cclip_duration = sum([c.duration for c in self.cclip])

                    # add new text clips
                    while self.cclip_duration < clip_duration:
                        cc = next(self.caption)
                        self.cclip.append(cc)
                        self.cclip_time = self.cclip[0].time
                        self.cclip_duration = sum([c.duration for c in self.cclip])

                    cc_text = " ".join([c.text for c in self.cclip])

                if self.dataset.text_transform:
                    cc_text = self.dataset.text_transform(cc_text)

                result[YouTubeDataset.F_CC_TEXT] = cc_text

            if YouTubeDataset.F_VF_DATA in self.dataset.fields:
                # get the video segment
                if self.dataset.video_cliplen is None:
                    if self.dataset.video_framerate == None:
                        video_cliplen = int(clip_duration * self.video_rate)
                    else:
                        video_cliplen = int(clip_duration * self.dataset.video_framerate)
                else:
                    video_cliplen = self.dataset.video_cliplen

                if self.dataset.video_stridelen is None:
                    stride = video_cliplen
                elif isinstance(self.dataset.video_stridelen, tuple):
                    stride = random.randint(*self.dataset.video_stridelen)
                else:
                    stride = self.dataset.video_stridelen

                try:
                    if self.dataset.video_framerate == None:
                        video_chunklen = max(stride,video_cliplen) + 1
                    else:
                        video_chunklen = max(stride,video_cliplen) * self.video_rate // self.dataset.video_framerate + 1

                    while len(self.vclip) < video_chunklen:
                        vf = next(self.video)
                        #vf_data = torch.from_numpy(np.transpose(vf.to_ndarray(format=self.dataset.video_format), (2, 0, 1))) # (C, H, W) Image Tensor
                        vf_data = torch.from_numpy(vf.to_ndarray(format=self.dataset.video_format)).permute(2, 0, 1) # (C, H, W) Image Tensor
                        if self.dataset.image_transform:
                            vf_data = self.dataset.image_transform(vf_data)
                        self.vclip.append((vf_data, vf.time))

                except StopIteration:
#                    if len(self.vclip) == 0:
                    raise StopIteration()

                if self.dataset.video_framerate != None and self.dataset.video_framerate != self.video_rate:
                    stride = stride * self.video_rate // self.dataset.video_framerate
                    num_frames = len(self.vclip) * self.dataset.video_framerate // self.video_rate
                    idx = YTDSIterator._resample_video_idx(num_frames, self.video_rate, self.dataset.video_framerate)
                    if isinstance(idx, slice):
                        vclip = self.vclip[idx]
                    else:
                        vclip = [ self.vclip[i] for i in idx ]
                else:
                    vclip = self.vclip

                #vf_data = torch.stack([ i[0] for i in vclip[:video_cliplen]]).permute(1, 0, 2, 3) # (C,T,H,W) - Video Tensor
                vf_data = torch.stack([ i[0] for i in vclip[:video_cliplen]]).permute(0, 2, 3, 1) # (T,H,W,C) - Video Tensor
                #vf_data = torch.stack([ i[0] for i in vclip[:video_cliplen]]) # (T,C,H,W) - Video Tensor

                if self.dataset.video_transform:
                    vf_data = self.dataset.video_transform(vf_data)
                result[YouTubeDataset.F_VF_DATA] = vf_data

                self.vclip_time = vclip[0][1]

                self.vclip = vclip[stride:]

                if self.dataset.key == YouTubeDataset.F_VF_DATA:
                    result[YouTubeDataset.F_TIME] = clip_time = self.vclip_time
                    result[YouTubeDataset.F_DURATION] = clip_duration

            if YouTubeDataset.F_IF_DATA in self.dataset.fields:

                if self.dataset.video_stridelen is None:
                    stride = 1
                elif isinstance(self.dataset.video_stridelen, tuple):
                    stride = random.randint(*self.dataset.video_stridelen)
                else:
                    stride = self.dataset.video_stridelen

                # get the video frame or segment
                for i in range(stride):
                    vf = next(self.video)
                vf_data = torch.from_numpy(np.transpose(vf.to_ndarray(format=self.dataset.video_format), (2, 0, 1))) # (C, H, W) Image Tensor
                if self.dataset.image_transform:
                    vf_data = self.dataset.image_transform(vf_data)
                result[YouTubeDataset.F_IF_DATA] = vf_data
                self.vclip_time = vf.time

                if self.dataset.key == YouTubeDataset.F_IF_DATA:
                    result[YouTubeDataset.F_TIME] = clip_time = self.vclip_time
                    result[YouTubeDataset.F_DURATION] = clip_duration

            if YouTubeDataset.F_AF_DATA in self.dataset.fields:

                if self.dataset.audio_cliplen is None:
                    audio_cliplen = int(clip_duration * self.dataset.audio_sr)
                else:
                    audio_cliplen = self.dataset.audio_cliplen

                # remove old audio clips
                for ix, (c, t)  in enumerate(self.aclip):
                    if (t + c.shape[-1] / self.dataset.audio_sr) < clip_time:
                        self.aclip.pop(ix)
                        if len(self.aclip) > 0:
                            self.aclip_time = self.aclip[0][1]
                        else:
                            self.aclip_time = 0

                # add new audio as needed
                while sum([c.shape[-1] for c, t in self.aclip]) < audio_cliplen:
                    af = next(self.audio)

                    if self.audio_resampler:
                        af.pts = None
                        af = self.audio_resampler.resample(af)

                    # get the audio frame as a numpy array and append to aclip
                    self.aclip.append((torch.from_numpy(af.to_ndarray(format=self.dataset.audio_format)), af.time))

                af_data = torch.cat([ c for c, t in self.aclip],dim=1) # (C,S) - Audio Tensor

                self.aclip = [ (af_data[:,:audio_cliplen], self.aclip[0][1])]
                # reshape aclip to make first segment audio_cliplen in size
                if af_data.shape[-1] > audio_cliplen:
                    self.aclip.append((af_data[:,audio_cliplen:], self.aclip_time + audio_cliplen / self.dataset.audio_sr))

                af_data = self.aclip[0][0]
                if self.dataset.audio_transform:
                    af_data = self.dataset.audio_transform(af_data)

                result[YouTubeDataset.F_AF_DATA] =  af_data
                if self.dataset.key == YouTubeDataset.F_AF_DATA:
                    result[YouTubeDataset.F_TIME] = self.aclip[0][1]
                    result[YouTubeDataset.F_DURATION] = clip_duration

            return [result[k] for k in self.dataset.fields]

        except StopIteration:
            self.next_active()
            return self.__next__()

    @staticmethod
    def _resample_video_idx(num_frames, original_fps, new_fps):
        step = float(original_fps) / new_fps
        if step.is_integer():
            # optimization: if step is integer, don't need to perform
            # advanced indexing
            step = int(step)
            return slice(None, None, step)
        idxs = torch.arange(num_frames, dtype=torch.float32) * step
        idxs = idxs.floor().to(torch.int64)
        return idxs

def YTDSVideoDownloadThunk(t):
    return t[0].download_video(t[1])

class YouTubeDatasetDownloader():
    def __init__(self, root_dir, channel, splits, api_key,
                 channel_id=None, user_name=None, caption_lang="en", types=['VIDEO', 'CAPTION'], youtube_url = 'http://youtube.com/watch?v='):

        assert channel_id != None or user_name != None, "channel_id or user_name must be specified"

        assert sum([splits[k] for k in splits]) == 1.0, "sum of splits must equal 1.0"

        self.root_dir = root_dir
        self.channel = channel
        self.ytapi = youtube.API(client_id=None, client_secret=None, api_key=api_key)
        self.channel_id = channel_id
        self.user_name = user_name
        self.splits = splits
        self.youtube_url = youtube_url
        self.caption_lang = caption_lang
        self.types=types

        self.channel_path = os.path.join(root_dir, channel)
        os.makedirs(self.channel_path, exist_ok=True)

        if 'VIDEO' in self.types:
            self.video_path = os.path.join(self.channel_path, 'video')
            os.makedirs(self.video_path, exist_ok=True)

        if 'CAPTION' in self.types:
            self.caption_path = os.path.join(self.channel_path, 'caption')
            os.makedirs(self.caption_path, exist_ok=True)

    def channel_ids(self):
        if self.channel_id != None:
            channels = self.ytapi.get('channels', part='snippet,contentDetails', id=self.channel_id)
        if self.user_name != None:
            channels = self.ytapi.get('channels', part='snippet,contentDetails', forUsername=self.user_name)
        return [ item['id'] for item in channels['items'] ]

    @staticmethod
    def playlist_sort_key(a):
        return a[1]

    def playlist_ids(self, channel_ids):
        playlist_ids = []
        for channel_id in channel_ids:

            nextPageToken=''

            while True:
                pl = self.ytapi.get('playlists', part='contentDetails', channelId=channel_id, pageToken=nextPageToken)

                for it in pl['items']:
                    if it['kind'] == 'youtube#playlist':
                        playlist_ids.append((it['id'], int(it['contentDetails']['itemCount'])))

                if 'nextPageToken' not in pl.keys():
                    break
                else:
                    nextPageToken=pl['nextPageToken']

        playlist_ids.sort(key=self.playlist_sort_key, reverse=True)

        return playlist_ids

    def videos(self, playlist_ids):
        videos = []

        for playlist_id, playlist_size in tqdm(playlist_ids):
            nextPageToken=''
            while True:
                pli = self.ytapi.get('playlistItems', part='snippet,contentDetails', playlistId=playlist_id, pageToken=nextPageToken)

                for it in pli['items']:
                    if it['kind'] == 'youtube#playlistItem':
                        snippet = it['snippet']
                        if snippet['resourceId']['kind'] == 'youtube#video':
                            videos.append([snippet['resourceId']['videoId'], snippet['title'], snippet['description']])

                if 'nextPageToken' not in pli.keys():
                    break
                else:
                    nextPageToken=pli['nextPageToken']

        return videos

    def make_dataframe(self, videos):
        self.df = pd.DataFrame(videos, columns =['video_id', 'title', 'description'], dtype = float)
        return self.df

    def parallel_download_videos(self, num_partitions=1, num_cores=1):
        df_split = np.array_split(self.df, num_partitions)
        selves = [self for _ in df_split]

        pool = Pool(num_cores)
        newdf = pd.concat(pool.map(YTDSVideoDownloadThunk, zip(selves, df_split)))
        pool.close()
        pool.join()
        self.df = newdf

    def download_videos(self, subframe):
        drops = []
        for ix, video_id in enumerate(subframe['video_id']):

            url = self.youtube_url + video_id
            try:
                yt = pytube.YouTube(url)

                if yt.captions[self.caption_lang] != None:
                    #TODO: Download at a specific resolution - currently highest
                    if 'VIDEO' in self.types:
                        stream = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').asc().first()
                        video_filename = video_id #+ '.mp4'
                        stream.download(output_path=self.video_path, filename=video_filename)

                    if 'CAPTION' in self.types:
                        caption = yt.captions[self.caption_lang]
                        caption_filename = os.path.join(self.caption_path, video_id + ".xml")
                        with open(caption_filename, "w") as caption_file:
                            caption_file.write(caption.xml_captions)
                else:
                    drops.append(ix)

            except OSError as ex:
                raise ex

            except Exception as ex:
                print('pytube3 incompatibility, upgrade pytube3 to download: ' + url, 'Exception: ' + str(ex.__class__.__name__) + ', '+ str(ex))
                drops.append(ix)

            self.pbar.update()

        subframe = subframe.drop(drops)
        return subframe

    def make_splits(self):
        sz = self.df.shape[0]
        perm = np.random.permutation(range(sz))

        hwm = 0
        for k in self.splits:
            start = hwm
            end = min(hwm + int(self.splits[k] * sz), sz)
            hwm = end + 1

            split = self.df.iloc[perm[start:end]]
            split.to_csv(os.path.join(self.channel_path, k + '.csv'))

    def splits_exist(self):
        return all([ os.path.isfile(os.path.join(self.channel_path, k + '.csv')) for k in self.splits ])

    def download_sync(self):
        if not self.splits_exist():
            channel_ids = self.channel_ids()
            playlist_ids = self.playlist_ids(channel_ids)
            videos = self.videos(playlist_ids)
            self.make_dataframe(videos)

            self.pbar = tqdm(total=self.df.shape[0])
            #self.parallel_download_videos()
            self.df = self.download_videos(self.df)
            self.pbar.close()
            self.make_splits()
        else:
            print('Files already downloaded')

max_stars_repo_name: johnrcumming/YouTubeDataset

max_issues_repo_name: johnrcumming/YouTubeDataset

max_forks_repo_name: johnrcumming/YouTubeDataset