edu.cuhk.hccl.cmd.AppQueryExpander.java Source code

Introduction

Here is the source code for edu.cuhk.hccl.cmd.AppQueryExpander.java

Source

/**
 * Copyright (C) 2014 Pengfei Liu <pfliu@se.cuhk.edu.hk>
 * The Chinese University of Hong Kong.
 *
 * This file is part of smart-search-web.
 *
 * 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 edu.cuhk.hccl.cmd;

import java.util.List;

import org.apache.commons.cli.BasicParser;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.CommandLineParser;
import org.apache.commons.cli.Options;
import org.apache.commons.cli.ParseException;

import edu.cuhk.hccl.expander.QueryExpander;
import edu.cuhk.hccl.expander.WordNetExpander;
import edu.cuhk.hccl.expander.WordVectorExpander;

/**
 * This program expands search query by finding similar terms via word2vec.
 * 
 * @author Pengfei Liu (pfliu@se.cuhk.edu.hk)
 * 
 */
public class AppQueryExpander {

    public static void main(String[] args) {

        try {

            CommandLineParser parser = new BasicParser();
            Options options = createOptions();
            CommandLine line = parser.parse(options, args);

            // Get parameters
            String queryStr = line.getOptionValue('q');
            int topK = Integer.parseInt(line.getOptionValue('k'));
            String modelPath = line.getOptionValue('m');
            String queryType = line.getOptionValue('t');

            QueryExpander expander = null;
            if (queryType.equalsIgnoreCase("WordVector")) {
                // Load word vectors
                System.out.println("[INFO] Loading word vectors...");
                expander = new WordVectorExpander(modelPath);
            } else if (queryType.equalsIgnoreCase("WordNet")) {
                // Load WordNet
                System.out.println("[INFO] Loading WordNet...");
                expander = new WordNetExpander(modelPath);
            } else {
                System.out.println("Please choose a correct expander: WordNet or WordVector!");
                System.exit(-1);
            }

            // Expand query
            System.out.println("[INFO] Computing similar words...");
            List<String> terms = expander.expandQuery(queryStr, topK);

            System.out.println(String.format("\n[INFO] The top %d similar words are: ", topK));
            for (String term : terms) {
                System.out.println(term);
            }

        } catch (ParseException exp) {
            System.out.println("Error in parameters: \n" + exp.getMessage());
            System.exit(-1);

        }
    }

    private static Options createOptions() {

        Options options = new Options();

        options.addOption("q", "query", true, "Query string");
        options.addOption("k", "topK", true, "Number of terms to return");
        options.addOption("m", "model", true, "Model file for word vectors or WordNet");
        options.addOption("t", "type", true, "Expand type: wordnet or wordvector");
        return options;
    }

}