eu.fbk.twm.wiki.xmldump.WikipediaTemplateContentExtractor Maven / Gradle / Ivy

Go to download
/*
 * Copyright (2013) Fondazione Bruno Kessler (http://www.fbk.eu/)
 *
 * 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.
 */

package eu.fbk.twm.wiki.xmldump;

import eu.fbk.twm.utils.CommandLineWithLogger;
import eu.fbk.twm.utils.ExtractorParameters;
import eu.fbk.twm.utils.GenericFileUtils;
import eu.fbk.twm.utils.WikipediaExtractor;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.OptionBuilder;
import org.apache.log4j.Logger;
import org.apache.log4j.PropertyConfigurator;
import org.apache.lucene.analysis.WhitespaceAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.store.FSDirectory;

import java.io.File;
import java.io.IOException;
import java.util.Locale;

public class WikipediaTemplateContentExtractor extends AbstractWikipediaExtractor implements WikipediaExtractor {
	/**
	 * Define a static logger variable so that it references the
	 * Logger instance named WikipediaTemplateExtractor.
	 */
	static Logger logger = Logger.getLogger(WikipediaTemplateContentExtractor.class.getName());

	private IndexWriter templateWriter;

	public WikipediaTemplateContentExtractor(int numThreads, int numPages, Locale locale) {
		super(numThreads, numPages, locale, "configuration/");
	}

	@Override
	public void start(ExtractorParameters extractorParameters) {
		try {
			String dirName = extractorParameters.getWikipediaTemplateFilePrefixName() + "xml-index-raw/";
			dirName = GenericFileUtils.checkWriteableFolder(dirName, true);
			templateWriter = new IndexWriter(FSDirectory.open(new File(dirName)), new WhitespaceAnalyzer(), IndexWriter.MaxFieldLength.LIMITED);

		} catch (IOException e) {
			e.printStackTrace();
			logger.error(e);
		}
		startProcess(extractorParameters.getWikipediaXmlFileName());
	}

	@Override
	public void disambiguationPage(String text, String title, int wikiID) {
		//To change body of implemented methods use File | Settings | File Templates.
	}

	@Override
	public void categoryPage(String text, String title, int wikiID) {
		//To change body of implemented methods use File | Settings | File Templates.
	}

	@Override
	public void redirectPage(String text, String title, int wikiID) {
		//To change body of implemented methods use File | Settings | File Templates.
	}

	@Override
	public void portalPage(String text, String title, int wikiID) {
		//To change body of implemented methods use File | Settings | File Templates.
	}

	@Override
	public void projectPage(String text, String title, int wikiID) {
		//To change body of implemented methods use File | Settings | File Templates.
	}

	@Override
	public void filePage(String text, String title, int wikiID) {
		//To change body of implemented methods use File | Settings | File Templates.
	}


	@Override
	public void templatePage(String text, String title, int wikiID) {
		try {
			Document doc = new Document();

			doc.add(new Field("page", title.toLowerCase(), Field.Store.YES, Field.Index.NOT_ANALYZED));
			doc.add(new Field("xml", text.getBytes(), Field.Store.YES));
			synchronized (this) {
				templateWriter.addDocument(doc);
			}
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	@Override
	public void contentPage(String text, String title, int wikiID) {
	}

	@Override
	public void endProcess() {
		super.endProcess();
		try {
			templateWriter.optimize();
			templateWriter.close();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public static void main(String[] args) {
		
		CommandLineWithLogger commandLineWithLogger = new CommandLineWithLogger();

		commandLineWithLogger.addOption(OptionBuilder.withLongOpt("in-wiki").withDescription("Input XML file").isRequired().hasArg().withArgName("filename").create("i"));
		commandLineWithLogger.addOption(OptionBuilder.withLongOpt("output").withDescription("Output base folder").isRequired().hasArg().withArgName("folder").create("o"));
		commandLineWithLogger.addOption(OptionBuilder.withLongOpt("threads").withDescription("Number of threads").hasArg().withArgName("num").create("t"));

		CommandLine commandLine = null;
		try {
			commandLine = commandLineWithLogger.getCommandLine(args);
			PropertyConfigurator.configure(commandLineWithLogger.getLoggerProps());
		} catch (Exception e) {
			System.exit(1);
		}
		
		String xin = commandLine.getOptionValue("in-wiki");
		String xout = commandLine.getOptionValue("output");

		Integer numThreads = 1;
		if (commandLine.hasOption('t')) {
			numThreads = Integer.parseInt(commandLine.getOptionValue('t'));
		}

		ExtractorParameters e = new ExtractorParameters(xin, xout, true);

		WikipediaTemplateContentExtractor writer = new WikipediaTemplateContentExtractor(numThreads, Integer.MAX_VALUE, e.getLocale());
		writer.start(e);
	}

}