<commit_before>package de.tudarmstadt.lt.structuredtopics.evaluate;

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.nio.charset.Charset;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.DefaultParser;
import org.apache.commons.cli.HelpFormatter;
import org.apache.commons.cli.Option;
import org.apache.commons.cli.Options;
import org.apache.commons.cli.ParseException;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.settings.Settings.Builder;
import org.elasticsearch.common.xcontent.XContentFactory;
import org.elasticsearch.node.Node;
import org.elasticsearch.node.NodeBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.google.common.io.Files;

import de.tudarmstadt.lt.structuredtopics.Main.InputMode;
import de.tudarmstadt.lt.structuredtopics.Utils;

public class Indexer {

	public static final String EL_FIELD_CLUSTER_LABEL = "label";
	public static final String EL_FIELD_CLUSTER_WORDS = "words";
	public static final String EL_INDEX_TYPE = "cluster";
	public static final String EL_INDEX = "topics";

	private static final String OPTION_INPUT_CLUSTERS = "clusters";
	private static final String OPTION_INPUT_LABELS = "labels";
	private static final String OPTION_INDEX_DIR = "index";

	private static final Logger LOG = LoggerFactory.getLogger(Indexer.class);

	public static void main(String[] args) {
		Options options = createOptions();
		Node node = null;
		try {
			CommandLine cl = new DefaultParser().parse(options, args, true);
			File clusters = new File(cl.getOptionValue(OPTION_INPUT_CLUSTERS));
			File clusterLabelsFile = new File(cl.getOptionValue(OPTION_INPUT_LABELS));
			Map<String, String> clusterLabels = readWordLabels(clusterLabelsFile);
			Builder settings = Settings.settingsBuilder().put("path.home", cl.getOptionValue(OPTION_INDEX_DIR));
			node = NodeBuilder.nodeBuilder().clusterName("structured-topics-el").settings(settings).node();
			Client client = node.client();
			client.admin().cluster().prepareHealth().setWaitForYellowStatus().execute().actionGet();
			buildIndex(clusters, clusterLabels, client);
		} catch (ParseException e) {
			LOG.error("Invalid arguments: {}", e.getMessage());
			StringWriter sw = new StringWriter();
			try (PrintWriter w = new PrintWriter(sw)) {
				new HelpFormatter().printHelp(w, Integer.MAX_VALUE, "application", "", options, 0, 0, "", true);
			}
			LOG.error(sw.toString());
		} catch (Exception e) {
			LOG.error("Error", e);
		} finally {
			if (node != null) {
				node.close();
			}
		}
	}

	private static void buildIndex(File clusters, Map<String, String> wordLabels, Client client) throws IOException {
		try (BufferedReader in = Utils.openReader(clusters, InputMode.GZ)) {
			String line = null;
			while ((line = in.readLine()) != null) {
				String[] split = line.split("\\t");
				if (split.length != 3) {
					LOG.error("Invalid cluster: {}", line);
				}
				Integer clusterId = Integer.valueOf(split[0]);
				Integer clusterSize = Integer.valueOf(split[1]);
				String lowerCase = split[2].toLowerCase();
				String[] clusterWords = lowerCase.split(",\\s*");
				String[] labels = getLabels(clusterWords, wordLabels);
				removePostTagAndSenseId(clusterWords);
				if (labels.length > 0) {
					// clusters without labels are skipped
					IndexResponse response = client.prepareIndex(EL_INDEX, EL_INDEX_TYPE)
							.setSource(XContentFactory.jsonBuilder().startObject()
									.field(EL_FIELD_CLUSTER_WORDS, clusterWords).field(EL_FIELD_CLUSTER_LABEL, labels)
									.endObject())
							.get();
					LOG.info(response.toString());
				}
			}
		}
	}

	private static String[] getLabels(String[] words, Map<String, String> wordLabels) {
		Set<String> labels = Sets.newHashSet();
		for (String word : words) {
			String label = wordLabels.get(word);
			if (label != null) {
				labels.add(label);
			}
		}
		return labels.toArray(new String[0]);
	}

	private static void removePostTagAndSenseId(String[] clusterWords) {
		for (int i = 0; i < clusterWords.length; i++) {
			String fullWord = clusterWords[i];
			int firstHashIndex = fullWord.indexOf("#");
			if (firstHashIndex != -1) {
				String word = fullWord.substring(0, firstHashIndex);
				clusterWords[i] = word;
			}
		}

	}

	private static Map<String, String> readWordLabels(File labelsFile) {
		Map<String, String> wordLabels = Maps.newHashMap();
		try {
			List<String> lines = Files.readLines(labelsFile, Charset.defaultCharset());
			for (String line : lines) {
				String[] split = line.toLowerCase().split("\t");
				if (split.length < 2) {
					continue;
				}
				String word = split[0];
				String label = split[1];
				wordLabels.put(word, label);
			}
		} catch (IOException e) {
			LOG.error("Error while reading labels", e);
		}
		return wordLabels;
	}

	private static Options createOptions() {
		Options options = new Options();
		Option inputClusters = Option.builder(OPTION_INPUT_CLUSTERS).argName("sense clusters")
				.desc("Input file containing the clusters, in .gz-format").hasArg().required().build();
		options.addOption(inputClusters);
		Option inputLabels = Option.builder(OPTION_INPUT_LABELS).argName("cluster labels")
				.desc(".csv file, containing a word#postag#senseid, eg. java#NN#1 and any label, separated by tab")
				.hasArg().required().build();
		options.addOption(inputLabels);
		Option indexDir = Option.builder(OPTION_INDEX_DIR).argName("index directory")
				.desc("path to directory, where the elasticsearch-index will be created").hasArg().required().build();
		options.addOption(indexDir);<commit_msg>[eval] write clusters with labels on indexing<commit_after>package de.tudarmstadt.lt.structuredtopics.evaluate;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.nio.charset.Charset;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.DefaultParser;
import org.apache.commons.cli.HelpFormatter;
import org.apache.commons.cli.Option;
import org.apache.commons.cli.Options;
import org.apache.commons.cli.ParseException;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.settings.Settings.Builder;
import org.elasticsearch.common.xcontent.XContentFactory;
import org.elasticsearch.node.Node;
import org.elasticsearch.node.NodeBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.Maps;
import com.google.common.collect.Sets;
import com.google.common.io.Files;

import de.tudarmstadt.lt.structuredtopics.Main.InputMode;
import de.tudarmstadt.lt.structuredtopics.Utils;

public class Indexer {

	public static final String EL_FIELD_CLUSTER_LABEL = "label";
	public static final String EL_FIELD_CLUSTER_WORDS = "words";
	public static final String EL_INDEX_TYPE = "cluster";
	public static final String EL_INDEX = "topics";

	private static final String OPTION_INPUT_CLUSTERS = "clusters";
	private static final String OPTION_INPUT_LABELS = "labels";
	private static final String OPTION_INDEX_DIR = "index";
	private static final String OPTION_LABELED_CLUSTERS = "labeledClustersOut";

	private static final Logger LOG = LoggerFactory.getLogger(Indexer.class);

	public static void main(String[] args) {
		Options options = createOptions();
		Node node = null;
		try {
			CommandLine cl = new DefaultParser().parse(options, args, true);
			File clusters = new File(cl.getOptionValue(OPTION_INPUT_CLUSTERS));
			File clusterLabelsFile = new File(cl.getOptionValue(OPTION_INPUT_LABELS));
			File labeledClustersOutFile = new File(cl.getOptionValue(OPTION_LABELED_CLUSTERS));
			Map<String, String> clusterLabels = readWordLabels(clusterLabelsFile);
			Builder settings = Settings.settingsBuilder().put("path.home", cl.getOptionValue(OPTION_INDEX_DIR));
			node = NodeBuilder.nodeBuilder().clusterName("structured-topics-el").settings(settings).node();
			Client client = node.client();
			client.admin().cluster().prepareHealth().setWaitForYellowStatus().execute().actionGet();
			buildIndex(clusters, clusterLabels, client, labeledClustersOutFile);
		} catch (ParseException e) {
			LOG.error("Invalid arguments: {}", e.getMessage());
			StringWriter sw = new StringWriter();
			try (PrintWriter w = new PrintWriter(sw)) {
				new HelpFormatter().printHelp(w, Integer.MAX_VALUE, "application", "", options, 0, 0, "", true);
			}
			LOG.error(sw.toString());
		} catch (Exception e) {
			LOG.error("Error", e);
		} finally {
			if (node != null) {
				node.close();
			}
		}
	}

	private static void buildIndex(File clusters, Map<String, String> wordLabels, Client client,
			File labeledClustersOutFile) throws IOException {
		try (BufferedWriter out = Utils.openGzipWriter(labeledClustersOutFile)) {
			try (BufferedReader in = Utils.openReader(clusters, InputMode.GZ)) {
				String line = null;
				while ((line = in.readLine()) != null) {
					String[] split = line.split("\\t");
					if (split.length != 3) {
						LOG.error("Invalid cluster: {}", line);
					}
					Integer clusterId = Integer.valueOf(split[0]);
					Integer clusterSize = Integer.valueOf(split[1]);
					String clusterWordsPlain = split[2].toLowerCase();
					String[] clusterWords = clusterWordsPlain.split(",\\s*");
					String[] labels = getLabels(clusterWords, wordLabels);
					writeLabeledCluster(out, clusterId, clusterSize, labels, clusterWordsPlain);
					removePostTagAndSenseId(clusterWords);
					if (labels.length > 0) {
						// clusters without labels are skipped
						IndexResponse response = client.prepareIndex(EL_INDEX, EL_INDEX_TYPE)
								.setSource(XContentFactory.jsonBuilder().startObject()
										.field(EL_FIELD_CLUSTER_WORDS, clusterWords)
										.field(EL_FIELD_CLUSTER_LABEL, labels).endObject())
								.get();
						LOG.info(response.toString());
					}
				}
			}
		}
	}

	private static void writeLabeledCluster(BufferedWriter out, Integer clusterId, Integer clusterSize, String[] labels,
			String clusterWordsPlain) {
		StringBuilder b = new StringBuilder();
		b.append(clusterId).append("\t").append(clusterSize).append("\t");
		for (String label : labels) {
			b.append(label).append(", ");
		}
		b.append("\t").append(clusterWordsPlain).append("\n");
		try {
			out.write(b.toString());
		} catch (IOException e) {
			LOG.error("Writing cluster with id {} failed", clusterId, e);
		}
	}

	private static String[] getLabels(String[] words, Map<String, String> wordLabels) {
		Set<String> labels = Sets.newHashSet();
		for (String word : words) {
			String label = wordLabels.get(word);
			if (label != null) {
				labels.add(label);
			}
		}
		return labels.toArray(new String[0]);
	}

	private static void removePostTagAndSenseId(String[] clusterWords) {
		for (int i = 0; i < clusterWords.length; i++) {
			String fullWord = clusterWords[i];
			int firstHashIndex = fullWord.indexOf("#");
			if (firstHashIndex != -1) {
				String word = fullWord.substring(0, firstHashIndex);
				clusterWords[i] = word;
			}
		}

	}

	private static Map<String, String> readWordLabels(File labelsFile) {
		Map<String, String> wordLabels = Maps.newHashMap();
		try {
			List<String> lines = Files.readLines(labelsFile, Charset.defaultCharset());
			for (String line : lines) {
				String[] split = line.toLowerCase().split("\t");
				if (split.length < 2) {
					continue;
				}
				String word = split[0];
				String label = split[1];
				wordLabels.put(word, label);
			}
		} catch (IOException e) {
			LOG.error("Error while reading labels", e);
		}
		return wordLabels;
	}

	private static Options createOptions() {
		Options options = new Options();
		Option inputClusters = Option.builder(OPTION_INPUT_CLUSTERS).argName("sense clusters")
				.desc("Input file containing the clusters, in .gz-format").hasArg().required().build();
		options.addOption(inputClusters);
		Option inputLabels = Option.builder(OPTION_INPUT_LABELS).argName("cluster labels")
				.desc(".csv file, containing a word#postag#senseid, eg. java#NN#1 and any label, separated by tab")
				.hasArg().required().build();
		options.addOption(inputLabels);
		Option indexDir = Option.builder(OPTION_INDEX_DIR).argName("index directory")
				.desc("path to directory, where the elasticsearch-index will be created").hasArg().required().build();
		options.addOption(indexDir);
		Option labeledClusters = Option.builder(OPTION_LABELED_CLUSTERS).argName("clusters with labels")
				.desc("new file with the same content as the clusters, but with an additional column for the labels of each cluster. Will be a csv.gz format, same as the clusters")
				.hasArg().required().build();
		options.addOption(labeledClusters);
{'fraction_non_alphanumeric': 0.12871961838419022, 'fraction_numerical': 0.002574392367683804, 'mean_word_length': 4.646857631466439, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 19, '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': '4158968', 'n_tokens_mistral': 4491, 'n_tokens_neox': 3980, 'n_words': 913}