<commit_before>/**
 * 
 */
package codemining.lm.grammar.tsg.tui;

import java.io.File;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.logging.Logger;

import org.apache.commons.io.FileUtils;
import org.apache.commons.io.filefilter.DirectoryFileFilter;
import org.apache.commons.io.filefilter.RegexFileFilter;
import org.apache.commons.lang.exception.ExceptionUtils;

import codemining.java.tokenizers.JavaTokenizer;
import codemining.lm.grammar.java.ast.BinaryJavaAstTreeExtractor;
import codemining.lm.grammar.java.ast.JavaAstTreeExtractor;
import codemining.lm.grammar.java.ast.VariableTypeJavaTreeExtractor;
import codemining.lm.grammar.tree.AbstractJavaTreeExtractor;
import codemining.lm.grammar.tree.TreeNode;
import codemining.lm.grammar.tsg.JavaFormattedTSGrammar;
import codemining.lm.grammar.tsg.TSGNode;
import codemining.lm.grammar.tsg.samplers.AbstractTSGSampler;
import codemining.lm.grammar.tsg.samplers.blocked.BlockCollapsedGibbsSampler;
import codemining.lm.grammar.tsg.samplers.blocked.FilteredBlockCollapsedGibbsSampler;
import codemining.lm.grammar.tsg.samplers.blocked.TreeCorpusFilter;
import codemining.util.SettingsLoader;
import codemining.util.serialization.ISerializationStrategy.SerializationException;
import codemining.util.serialization.Serializer;

/**
 * Sample a TSG using a blocked sampler.
 * 
 * @author <NAME> <<EMAIL>>
 * 
 */
public class SampleBlockedTSG {

	private static final int TREE_SPLIT_CFG_COUNT = (int) SettingsLoader
			.getNumericSetting("treeSplitCfgCount", 10);
	private static final Logger LOGGER = Logger
			.getLogger(SampleBlockedTSG.class.getName());

	/**
	 * @param args
	 * @throws SerializationException
	 */
	public static void main(final String[] args) throws SerializationException {
		if (args.length < 5) {
			System.err
					.println("Usage <TsgTrainingDir> normal|binary|binaryvariables|variables|binaryvariablesNoAnnotate block|filterblock <alpha> <#iterations> [<CfgExtraTraining>]");
			System.exit(-1);
		}

		final int nIterations = Integer.parseInt(args[4]);
		final double concentrationParameter = Double.parseDouble(args[3]);
		final File samplerCheckpoint = new File("tsgSampler.ser");
		final BlockCollapsedGibbsSampler sampler;

		if (samplerCheckpoint.exists()) {
			sampler = (BlockCollapsedGibbsSampler) Serializer.getSerializer()
					.deserializeFrom("tsgSampler.ser");
			LOGGER.info("Resuming sampling");

		} else {

			final AbstractJavaTreeExtractor format;
			if (args[1].equals("normal")) {
				format = new JavaAstTreeExtractor();
			} else if (args[1].equals("binary")) {
				format = new BinaryJavaAstTreeExtractor(
						new JavaAstTreeExtractor());
			} else if (args[1].equals("variables")) {
				format = new VariableTypeJavaTreeExtractor();
			} else if (args[1].equals("binaryvariables")) {
				format = new BinaryJavaAstTreeExtractor(
						new VariableTypeJavaTreeExtractor());
			} else if (args[1].equals("binaryvariablesNoAnnotate")) {
				format = new BinaryJavaAstTreeExtractor(
						new VariableTypeJavaTreeExtractor(), false);
			} else {
				throw new IllegalArgumentException(
						"Unrecognizable training type parameter " + args[1]);
			}

			if (args[2].equals("block")) {
				sampler = new BlockCollapsedGibbsSampler(100,
						concentrationParameter, new JavaFormattedTSGrammar(
								format), new JavaFormattedTSGrammar(format));
			} else if (args[2].equals("filterblock")) {
				sampler = new FilteredBlockCollapsedGibbsSampler(100,
						concentrationParameter, new JavaFormattedTSGrammar(
								format), new JavaFormattedTSGrammar(format));
			} else {
				throw new IllegalArgumentException(
						"Unrecognizable training type parameter " + args[2]);
			}

			if (args.length > 5) {
				LOGGER.info("Loading additional CFG prior information from "
						+ args[5]);
				for (final File fi : FileUtils.listFiles(new File(args[5]),
						new RegexFileFilter(".*\\.java$"),
						DirectoryFileFilter.DIRECTORY)) {
					try {
						final TreeNode<TSGNode> ast = TSGNode.convertTree(
								format.getTree(fi), 0);
						sampler.addDataToPrior(ast);
					} catch (final Exception e) {
						LOGGER.warning("Failed to get AST for Cfg Prior "
								+ fi.getAbsolutePath() + " "
								+ ExceptionUtils.getFullStackTrace(e));
					}
				}
			}

			final double percentRootsInit = .7;
			int nFiles = 0;
			int nNodes = 0;
			LOGGER.info("Loading sample trees from  " + args[0]);
			final TreeCorpusFilter filter = new TreeCorpusFilter(format,
					TREE_SPLIT_CFG_COUNT);
			for (final File fi : FileUtils.listFiles(new File(args[0]),
					JavaTokenizer.javaCodeFileFilter,
					DirectoryFileFilter.DIRECTORY)) {
				try {
					final TreeNode<TSGNode> ast = TSGNode.convertTree(
							format.getTree(fi), percentRootsInit);
					nNodes += ast.getTreeSize();
					nFiles++;
					filter.addTree(ast);
				} catch (final Exception e) {
					LOGGER.warning("Failed to get AST for "
							+ fi.getAbsolutePath() + " "
							+ ExceptionUtils.getFullStackTrace(e));
				}
			}
			LOGGER.info("Loaded " + nFiles + " files containing " + nNodes
					+ " nodes");
			for (final TreeNode<TSGNode> filteredTree : filter
					.getFilteredTrees()) {
				sampler.addTree(filteredTree);
			}
			sampler.lockSamplerData();
		}

		final AtomicBoolean finished = new AtomicBoolean(false);
		Runtime.getRuntime().addShutdownHook(new Thread() {
			@Override
			public void run() {
				int i = 0;
				while (!finished.get() && i < 1000) {
					try {
						Thread.sleep(500);<commit_msg>Mine only snippets that contain usages of the relevant class, ignoring the rest.<commit_after>/**
 * 
 */
package codemining.lm.grammar.tsg.tui;

import java.io.File;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.logging.Logger;

import org.apache.commons.io.FileUtils;
import org.apache.commons.io.filefilter.DirectoryFileFilter;
import org.apache.commons.io.filefilter.RegexFileFilter;
import org.apache.commons.lang.exception.ExceptionUtils;
import org.eclipse.jdt.core.dom.ASTNode;

import codemining.java.codeutils.UsagePointExtractor;
import codemining.java.tokenizers.JavaTokenizer;
import codemining.lm.grammar.java.ast.BinaryJavaAstTreeExtractor;
import codemining.lm.grammar.java.ast.JavaAstTreeExtractor;
import codemining.lm.grammar.java.ast.VariableTypeJavaTreeExtractor;
import codemining.lm.grammar.tree.AbstractJavaTreeExtractor;
import codemining.lm.grammar.tree.TreeNode;
import codemining.lm.grammar.tsg.JavaFormattedTSGrammar;
import codemining.lm.grammar.tsg.TSGNode;
import codemining.lm.grammar.tsg.samplers.AbstractTSGSampler;
import codemining.lm.grammar.tsg.samplers.blocked.BlockCollapsedGibbsSampler;
import codemining.lm.grammar.tsg.samplers.blocked.FilteredBlockCollapsedGibbsSampler;
import codemining.lm.grammar.tsg.samplers.blocked.TreeCorpusFilter;
import codemining.util.SettingsLoader;
import codemining.util.serialization.ISerializationStrategy.SerializationException;
import codemining.util.serialization.Serializer;

/**
 * Sample a TSG using a blocked sampler.
 * 
 * @author <NAME> <<EMAIL>>
 * 
 */
public class SampleBlockedTSG {

	private static final int TREE_SPLIT_CFG_COUNT = (int) SettingsLoader
			.getNumericSetting("treeSplitCfgCount", 0);

	/**
	 * Filter only the usages of a package.
	 */
	private static final String FILTER_ONLY_PACKAGE_USAGES = SettingsLoader
			.getStringSetting("classFilter", null);
	private static final Logger LOGGER = Logger
			.getLogger(SampleBlockedTSG.class.getName());

	/**
	 * @param args
	 * @throws SerializationException
	 */
	public static void main(final String[] args) throws SerializationException {
		if (args.length < 5) {
			System.err
					.println("Usage <TsgTrainingDir> normal|binary|binaryvariables|variables|binaryvariablesNoAnnotate block|filterblock <alpha> <#iterations> [<CfgExtraTraining>]");
			System.exit(-1);
		}

		final int nIterations = Integer.parseInt(args[4]);
		final double concentrationParameter = Double.parseDouble(args[3]);
		final File samplerCheckpoint = new File("tsgSampler.ser");
		final BlockCollapsedGibbsSampler sampler;

		if (samplerCheckpoint.exists()) {
			sampler = (BlockCollapsedGibbsSampler) Serializer.getSerializer()
					.deserializeFrom("tsgSampler.ser");
			LOGGER.info("Resuming sampling");

		} else {

			final AbstractJavaTreeExtractor format;
			if (args[1].equals("normal")) {
				format = new JavaAstTreeExtractor();
			} else if (args[1].equals("binary")) {
				format = new BinaryJavaAstTreeExtractor(
						new JavaAstTreeExtractor());
			} else if (args[1].equals("variables")) {
				format = new VariableTypeJavaTreeExtractor();
			} else if (args[1].equals("binaryvariables")) {
				format = new BinaryJavaAstTreeExtractor(
						new VariableTypeJavaTreeExtractor());
			} else if (args[1].equals("binaryvariablesNoAnnotate")) {
				format = new BinaryJavaAstTreeExtractor(
						new VariableTypeJavaTreeExtractor(), false);
			} else {
				throw new IllegalArgumentException(
						"Unrecognizable training type parameter " + args[1]);
			}

			if (args[2].equals("block")) {
				sampler = new BlockCollapsedGibbsSampler(100,
						concentrationParameter, new JavaFormattedTSGrammar(
								format), new JavaFormattedTSGrammar(format));
			} else if (args[2].equals("filterblock")) {
				sampler = new FilteredBlockCollapsedGibbsSampler(100,
						concentrationParameter, new JavaFormattedTSGrammar(
								format), new JavaFormattedTSGrammar(format));
			} else {
				throw new IllegalArgumentException(
						"Unrecognizable training type parameter " + args[2]);
			}

			if (args.length > 5) {
				LOGGER.info("Loading additional CFG prior information from "
						+ args[5]);
				for (final File fi : FileUtils.listFiles(new File(args[5]),
						new RegexFileFilter(".*\\.java$"),
						DirectoryFileFilter.DIRECTORY)) {
					try {
						final TreeNode<TSGNode> ast = TSGNode.convertTree(
								format.getTree(fi), 0);
						sampler.addDataToPrior(ast);
					} catch (final Exception e) {
						LOGGER.warning("Failed to get AST for Cfg Prior "
								+ fi.getAbsolutePath() + " "
								+ ExceptionUtils.getFullStackTrace(e));
					}
				}
			}

			final double percentRootsInit = .7;
			int nFiles = 0;
			int nNodes = 0;
			LOGGER.info("Loading sample trees from  " + args[0]);
			final TreeCorpusFilter filter = new TreeCorpusFilter(format,
					TREE_SPLIT_CFG_COUNT);
			for (final File fi : FileUtils.listFiles(new File(args[0]),
					JavaTokenizer.javaCodeFileFilter,
					DirectoryFileFilter.DIRECTORY)) {
				try {
					if (FILTER_ONLY_PACKAGE_USAGES == null) {
						final TreeNode<TSGNode> ast = TSGNode.convertTree(
								format.getTree(fi), percentRootsInit);
						nNodes += ast.getTreeSize();
						filter.addTree(ast);
					} else {
						for (final ASTNode tree : UsagePointExtractor
								.usagePoints(FILTER_ONLY_PACKAGE_USAGES, fi)) {
							final TreeNode<TSGNode> ast = TSGNode.convertTree(
									format.getTree(tree), percentRootsInit);
							nNodes += ast.getTreeSize();
							filter.addTree(ast);
						}
					}
					nFiles++;
				} catch (final Exception e) {
					LOGGER.warning("Failed to get AST for "
							+ fi.getAbsolutePath() + " "
							+ ExceptionUtils.getFullStackTrace(e));
				}
			}
			LOGGER.info("Loaded " + nFiles + " files containing " + nNodes
					+ " nodes");
			for (final TreeNode<TSGNode> filteredTree : filter
					.getFilteredTrees()) {
				sampler.addTree(filteredTree);
			}
			sampler.lockSamplerData();
		}

		final AtomicBoolean finished = new AtomicBoolean(false);
		Runtime.getRuntime().addShutdownHook(new Thread() {
			@Override
			public void run() {
				int i = 0;
				while (!finished.get() && i < 1000) {
					try {
						Thread.sleep(500);
{'fraction_non_alphanumeric': 0.11458066717856838, 'fraction_numerical': 0.0063987714358843106, 'mean_word_length': 4.1344721857205435, 'pattern_counts': {'":': 0, '<': 28, '<?xml version=': 0, '>': 26, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3019903', 'n_tokens_mistral': 4446, 'n_tokens_neox': 3627, 'n_words': 710}