<commit_before>package cc.mallet.topics;

import cc.mallet.util.Randoms;
import cc.mallet.util.CommandOption;
import cc.mallet.types.*;
import java.util.*;
import java.io.*;
import java.util.concurrent.*;

public class WordEmbeddings {
	
	static CommandOption.String inputFile = new CommandOption.String(WordEmbeddings.class, "input", "FILENAME", true, null,
																	 "The filename from which to read the list of training instances.  Use - for stdin.  " +
																	 "The instances must be FeatureSequence or FeatureSequenceWithBigrams, not FeatureVector", null);
	
	static CommandOption.String outputFile = new CommandOption.String(WordEmbeddings.class, "output", "FILENAME", true, "weights.txt",
																	  "The filename to write text-formatted word vectors.", null);
	
	static CommandOption.Integer numDimensions = new CommandOption.Integer(WordEmbeddings.class, "num-dimensions", "INTEGER", true, 50,
																	   "The number of dimensions to fit.", null);

	static CommandOption.Integer windowSizeOption = new CommandOption.Integer(WordEmbeddings.class, "window-size", "INTEGER", true, 5,
																		"The number of adjacent words to consider.", null);


	static CommandOption.Integer numThreads = new CommandOption.Integer(WordEmbeddings.class, "num-threads", "INTEGER", true, 1,
																		"The number of threads for parallel training.", null);

	static CommandOption.Integer numIterationsOption = new CommandOption.Integer(WordEmbeddings.class, "num-iters", "INTEGER", true, 3,
																		"The number of passes through the training data.", null);

	static CommandOption.Double samplingFactorOption = new CommandOption.Double(WordEmbeddings.class, "frequency-factor", "NUMBER", true, 0.0001,
																		"Down-sample words that account for more than ~2.5x this proportion or the corpus.", null);

	static CommandOption.Integer numSamples = new CommandOption.Integer(WordEmbeddings.class, "num-samples", "INTEGER", true, 5,
																		"The number of negative samples to use in training.", null);

	static CommandOption.String exampleWord = new CommandOption.String(WordEmbeddings.class, "example-word", "STRING", true, null,
																	   "If defined, periodically show the closest vectors to this word.", null);


	Alphabet vocabulary;
        
	int numWords;
	int numColumns;
	double[] weights;
	double[] negativeWeights;
	int stride;
	
	int numIterations;

	int[] wordCounts;
	double[] retentionProbability;
	double[] samplingDistribution;
	int[] samplingTable;
	int samplingTableSize = 100000000;
	double samplingSum = 0.0;
	int totalWords = 0;

	double maxExpValue = 6.0;
	double minExpValue = -6.0;
	double[] sigmoidCache;
	int sigmoidCacheSize = 1000;

	int windowSize = 5;

	public int getMinDocumentLength() {
		return minDocumentLength;
	}

	public void setMinDocumentLength(int minDocumentLength) {
		if (minDocumentLength <= 0) {
			throw new IllegalArgumentException("Minimum document length must be at least 1.");
		}
		this.minDocumentLength = minDocumentLength;
	}

	private int minDocumentLength = 10;

	public void setNumIterations(int i) { numIterations = i; }

	public String getQueryWord() {
		return queryWord;
	}

	public void setQueryWord(String queryWord) {
		this.queryWord = queryWord;
	}

	String queryWord = "the";

	Randoms random = new Randoms();

	public WordEmbeddings() { }

	public WordEmbeddings(Alphabet a, int numColumns, int windowSize) {
		vocabulary = a;
                
		numWords = vocabulary.size();

		System.out.format("Vocab size: %d\n", numWords);

		this.numColumns = numColumns;
		
		this.stride = numColumns;
		
		weights = new double[numWords * stride];
		negativeWeights = new double[numWords * stride];

		for (int word = 0; word < numWords; word++) {
			for (int col = 0; col < numColumns; col++) {
				weights[word * stride + col] = (random.nextDouble() - 0.5f) / numColumns;
				negativeWeights[word * stride + col] = 0.0;
			}
		}

		wordCounts = new int[numWords];
		samplingDistribution = new double[numWords];
		retentionProbability = new double[numWords];
		samplingTable = new int[samplingTableSize];

		this.windowSize = windowSize;

		sigmoidCache = new double[sigmoidCacheSize + 1];

		for (int i = 0; i < sigmoidCacheSize; i++) {
			double value = ((double) i / sigmoidCacheSize) * (maxExpValue - minExpValue) + minExpValue;
			sigmoidCache[i] = 1.0 / (1.0 + Math.exp(-value));
		}
		
	}

	public void countWords(InstanceList instances, double samplingFactor) {
		for (Instance instance: instances) {
			
			FeatureSequence tokens = (FeatureSequence) instance.getData();
			int length = tokens.getLength();

			for (int position = 0; position < length; position++) {
				int type = tokens.getIndexAtPosition(position);
				wordCounts[type]++;
			}
			
			totalWords += length;
		}
		
		for (int word = 0; word < numWords; word++) {
			// Word2Vec style sampling weight
			double frequencyScore = (double) wordCounts[word] / (samplingFactor * totalWords);
			retentionProbability[word] = Math.min((Math.sqrt(frequencyScore) + 1) / frequencyScore, 1.0);
		}

		IDSorter[] sortedWords = new IDSorter[numWords];
		for (int word = 0; word < numWords; word++) {
			sortedWords[word] = new IDSorter(word, wordCounts[word]);
		}
		Arrays.sort(sortedWords);
		
		samplingDistribution[0] = Math.pow(wordCounts[ sortedWords[0].getID() ], 0.75);
		for (int word = 1; word < numWords; word++) {
			samplingDistribution[word] = samplingDistribution[word-1] + Math.pow(wordCounts[ sortedWords[word].getID() ], 0.75);
		}
		samplingSum = samplingDistribution[numWords-1];

		int order = 0;
		for (int i = 0; i < samplingTableSize; i++) {
			samplingTable[i] = sortedWords[order].getID();
			while (samplingSum * i / samplingTableSize > samplingDistribution[order]) {
				order++;
			}
		}

		System.out.println("done counting: " + totalWords);
	}

	public void train(InstanceList instances, int numThreads, int numSamples) {

		ExecutorService executor = Executors.newFixedThreadPool(numThreads);

		WordEmbeddingRunnable[] runnables = new WordEmbeddingRunnable[numThreads];
		for (int thread = 0; thread < numThreads; thread++) {
			runnables[thread] = new WordEmbeddingRunnable(this, instances, numSamples, numThreads, thread);
			executor.submit(runnables[thread]);
		}

		long startTime = System.currentTimeMillis();
		double difference = 0.0;

		boolean finished = false;
		while (! finished) {
                               
			try {
				Thread.sleep(5000);
			} catch (InterruptedException e) {
                                                
			}

			int wordsSoFar = 0;
                                        
			// Are all the threads done?
			for (int thread = 0; thread < numThreads; thread++) {
				wordsSoFar += runnables[thread].wordsSoFar;
				//System.out.format("%.3f ", runnables[thread].getMeanError());
			}

			long runningMillis = System.currentTimeMillis() - startTime;
			System.out.format("%d\t%d\t%fk w/s %f avg\n", wordsSoFar, runningMillis, (double) wordsSoFar / runningMillis, 
							  averageAbsWeight());
			//variances();
			difference = 0.0;

			if (wordsSoFar > numIterations * totalWords) {
				finished = true;
				for (int thread = 0; thread < numThreads; thread++) {
					runnables[thread].shouldRun = false;
				}
			}

			if (queryWord != null && vocabulary.contains(queryWord)) {
				findClosest(copy(queryWord));
			}
		}
		executor.shutdownNow();
	}

	public void findClosest(double[] targetVector) {
		IDSorter[] sortedWords = new IDSorter[numWords];

		double targetSquaredSum = 0.0;
		for (int col = 0; col < numColumns; col++) {
			targetSquaredSum += targetVector[col] * targetVector[col];
		}
		double targetNormalizer = 1.0 / Math.sqrt(targetSquaredSum);

		System.out.println(targetSquaredSum);

		for (int word = 0; word < numWords; word++) {
			
			double innerProduct = 0.0;
			
			double wordSquaredSum = 0.0;
			for (int col = 0; col < numColumns; col++) {
				wordSquaredSum += weights[word * stride + col] * weights[word * stride + col];
			}
			double wordNormalizer = 1.0 / Math.sqrt(wordSquaredSum);

			for (int col = 0; col < numColumns; col++) {
				innerProduct += targetNormalizer * targetVector[col] * wordNormalizer * weights[word * stride + col];
			}

			sortedWords[word] = new IDSorter(word, innerProduct);
		}

		Arrays.sort(sortedWords);
		
		for (int i = 0; i < 10; i++) {
			System.out.format("%f\t%d\t%s\n", sortedWords[i].getWeight(), sortedWords[i].getID(), vocabulary.lookupObject(sortedWords[i].getID()));
		}
	}
	<commit_msg>reduced memory usage for diagnostics<commit_after>package cc.mallet.topics;

import cc.mallet.util.Randoms;
import cc.mallet.util.CommandOption;
import cc.mallet.types.*;
import java.util.*;
import java.io.*;
import java.util.concurrent.*;

public class WordEmbeddings {
	
	static CommandOption.String inputFile = new CommandOption.String(WordEmbeddings.class, "input", "FILENAME", true, null,
																	 "The filename from which to read the list of training instances.  Use - for stdin.  " +
																	 "The instances must be FeatureSequence or FeatureSequenceWithBigrams, not FeatureVector", null);
	
	static CommandOption.String outputFile = new CommandOption.String(WordEmbeddings.class, "output", "FILENAME", true, "weights.txt",
																	  "The filename to write text-formatted word vectors.", null);
	
	static CommandOption.Integer numDimensions = new CommandOption.Integer(WordEmbeddings.class, "num-dimensions", "INTEGER", true, 50,
																	   "The number of dimensions to fit.", null);

	static CommandOption.Integer windowSizeOption = new CommandOption.Integer(WordEmbeddings.class, "window-size", "INTEGER", true, 5,
																		"The number of adjacent words to consider.", null);


	static CommandOption.Integer numThreads = new CommandOption.Integer(WordEmbeddings.class, "num-threads", "INTEGER", true, 1,
																		"The number of threads for parallel training.", null);

	static CommandOption.Integer numIterationsOption = new CommandOption.Integer(WordEmbeddings.class, "num-iters", "INTEGER", true, 3,
																		"The number of passes through the training data.", null);

	static CommandOption.Double samplingFactorOption = new CommandOption.Double(WordEmbeddings.class, "frequency-factor", "NUMBER", true, 0.0001,
																		"Down-sample words that account for more than ~2.5x this proportion or the corpus.", null);

	static CommandOption.Integer numSamples = new CommandOption.Integer(WordEmbeddings.class, "num-samples", "INTEGER", true, 5,
																		"The number of negative samples to use in training.", null);

	static CommandOption.String exampleWord = new CommandOption.String(WordEmbeddings.class, "example-word", "STRING", true, null,
																	   "If defined, periodically show the closest vectors to this word.", null);


	Alphabet vocabulary;
        
	int numWords;
	int numColumns;
	double[] weights;
	double[] negativeWeights;
	int stride;
	
	int numIterations;

	int[] wordCounts;
	double[] retentionProbability;
	double[] samplingDistribution;
	int[] samplingTable;
	int samplingTableSize = 100000000;
	double samplingSum = 0.0;
	int totalWords = 0;

	double maxExpValue = 6.0;
	double minExpValue = -6.0;
	double[] sigmoidCache;
	int sigmoidCacheSize = 1000;

	int windowSize = 5;
	
	IDSorter[] sortedWords = null;

	public int getMinDocumentLength() {
		return minDocumentLength;
	}

	public void setMinDocumentLength(int minDocumentLength) {
		if (minDocumentLength <= 0) {
			throw new IllegalArgumentException("Minimum document length must be at least 1.");
		}
		this.minDocumentLength = minDocumentLength;
	}

	private int minDocumentLength = 10;

	public void setNumIterations(int i) { numIterations = i; }

	public String getQueryWord() {
		return queryWord;
	}

	public void setQueryWord(String queryWord) {
		this.queryWord = queryWord;
	}

	String queryWord = "the";

	Randoms random = new Randoms();

	public WordEmbeddings() { }

	public WordEmbeddings(Alphabet a, int numColumns, int windowSize) {
		vocabulary = a;
                
		numWords = vocabulary.size();

		System.out.format("Vocab size: %d\n", numWords);

		this.numColumns = numColumns;
		
		this.stride = numColumns;
		
		weights = new double[numWords * stride];
		negativeWeights = new double[numWords * stride];

		for (int word = 0; word < numWords; word++) {
			for (int col = 0; col < numColumns; col++) {
				weights[word * stride + col] = (random.nextDouble() - 0.5f) / numColumns;
				negativeWeights[word * stride + col] = 0.0;
			}
		}

		wordCounts = new int[numWords];
		samplingDistribution = new double[numWords];
		retentionProbability = new double[numWords];
		samplingTable = new int[samplingTableSize];

		this.windowSize = windowSize;

		sigmoidCache = new double[sigmoidCacheSize + 1];

		for (int i = 0; i < sigmoidCacheSize; i++) {
			double value = ((double) i / sigmoidCacheSize) * (maxExpValue - minExpValue) + minExpValue;
			sigmoidCache[i] = 1.0 / (1.0 + Math.exp(-value));
		}
		
	}
	
	public void initializeSortables() {
		sortedWords = new IDSorter[numWords];
		for (int word = 0; word < numWords; word++) {
			sortedWords[word] = new IDSorter(word, 0.0);
		}
	}

	public void countWords(InstanceList instances, double samplingFactor) {
		for (Instance instance: instances) {
			
			FeatureSequence tokens = (FeatureSequence) instance.getData();
			int length = tokens.getLength();

			for (int position = 0; position < length; position++) {
				int type = tokens.getIndexAtPosition(position);
				wordCounts[type]++;
			}
			
			totalWords += length;
		}
		
		for (int word = 0; word < numWords; word++) {
			// Word2Vec style sampling weight
			double frequencyScore = (double) wordCounts[word] / (samplingFactor * totalWords);
			retentionProbability[word] = Math.min((Math.sqrt(frequencyScore) + 1) / frequencyScore, 1.0);
		}

		if (sortedWords == null) {
			initializeSortables();
		}
		for (int word = 0; word < numWords; word++) {
			sortedWords[word].set(word, wordCounts[word]);
		}
		Arrays.sort(sortedWords);
		
		samplingDistribution[0] = Math.pow(wordCounts[ sortedWords[0].getID() ], 0.75);
		for (int word = 1; word < numWords; word++) {
			samplingDistribution[word] = samplingDistribution[word-1] + Math.pow(wordCounts[ sortedWords[word].getID() ], 0.75);
		}
		samplingSum = samplingDistribution[numWords-1];

		int order = 0;
		for (int i = 0; i < samplingTableSize; i++) {
			samplingTable[i] = sortedWords[order].getID();
			while (samplingSum * i / samplingTableSize > samplingDistribution[order]) {
				order++;
			}
		}

		System.out.println("done counting: " + totalWords);
	}

	public void train(InstanceList instances, int numThreads, int numSamples) {

		ExecutorService executor = Executors.newFixedThreadPool(numThreads);

		WordEmbeddingRunnable[] runnables = new WordEmbeddingRunnable[numThreads];
		for (int thread = 0; thread < numThreads; thread++) {
			runnables[thread] = new WordEmbeddingRunnable(this, instances, numSamples, numThreads, thread);
			executor.submit(runnables[thread]);
		}

		long startTime = System.currentTimeMillis();
		double difference = 0.0;

		boolean finished = false;
		while (! finished) {
                               
			try {
				Thread.sleep(5000);
			} catch (InterruptedException e) {
                                                
			}

			int wordsSoFar = 0;
                                        
			// Are all the threads done?
			for (int thread = 0; thread < numThreads; thread++) {
				wordsSoFar += runnables[thread].wordsSoFar;
				//System.out.format("%.3f ", runnables[thread].getMeanError());
			}

			long runningMillis = System.currentTimeMillis() - startTime;
			System.out.format("%d\t%d\t%fk w/s %f avg\n", wordsSoFar, runningMillis, (double) wordsSoFar / runningMillis, 
							  averageAbsWeight());
			//variances();
			difference = 0.0;

			if (wordsSoFar > numIterations * totalWords) {
				finished = true;
				for (int thread = 0; thread < numThreads; thread++) {
					runnables[thread].shouldRun = false;
				}
			}

			if (queryWord != null && vocabulary.contains(queryWord)) {
				findClosest(copy(queryWord));
			}
		}
		executor.shutdownNow();
	}

	public void findClosest(double[] targetVector) {
		if (sortedWords == null) {
			initializeSortables();
		}

		double targetSquaredSum = 0.0;
		for (int col = 0; col < numColumns; col++) {
			targetSquaredSum += targetVector[col] * targetVector[col];
		}
		double targetNormalizer = 1.0 / Math.sqrt(targetSquaredSum);

		for (int word = 0; word < numWords; word++) {
			
			double innerProduct = 0.0;
			
			double wordSquaredSum = 0.0;
			for (int col = 0; col < numColumns; col++) {
				wordSquaredSum += weights[word * stride + col] * weights[word * stride + col];
			}
			double wordNormalizer = 1.0 / Math.sqrt(wordSquaredSum);

			for (int col = 0; col < numColumns; col++) {
				innerProduct += targetVector[col] * weights[word * stride + col];
			}
			innerProduct *= targetNormalizer * wordNormalizer;

			sortedWords[word].set(word, innerProduct);
		}

		Arrays.sort(sortedWords);
		
		for (int i = 0; i < 10; i++) {
			System.out.format("%f\t%d\t%s\n", sortedWords[i].getWeight(), sortedWords[i].getID(), vocabulary.lookupObject(sortedWords[i].getID()));
		}
	}
	
{'fraction_non_alphanumeric': 0.121631000578369, 'fraction_numerical': 0.01174089068825911, 'mean_word_length': 3.525255168803978, 'pattern_counts': {'":': 0, '<': 38, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4879274', 'n_tokens_mistral': 6418, 'n_tokens_neox': 5187, 'n_words': 1382}