/* * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You 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 org.apache.lucene.classification; import java.io.IOException; import java.util.Random; import org.apache.lucene.analysis.Analyzer; import org.apache.lucene.document.Document; import org.apache.lucene.document.Field; import org.apache.lucene.document.FieldType; import org.apache.lucene.document.TextField; import org.apache.lucene.index.IndexWriterConfig; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.RandomIndexWriter; import org.apache.lucene.search.Query; import org.apache.lucene.store.Directory; import org.apache.lucene.util.BytesRef; import org.apache.lucene.util.LuceneTestCase; import org.apache.lucene.util.TestUtil; import org.junit.After; import org.junit.Before; /** * Base class for testing {@link Classifier}s */ public abstract class ClassificationTestBase<T> extends LuceneTestCase { protected static final String POLITICS_INPUT = "Here are some interesting questions and answers about Mitt Romney.. " + "If you don't know the answer to the question about Mitt Romney, then simply click on the answer below the question section."; protected static final BytesRef POLITICS_RESULT = new BytesRef("politics"); protected static final String TECHNOLOGY_INPUT = "Much is made of what the likes of Facebook, Google and Apple know about users." + " Truth is, Amazon may know more."; protected static final String STRONG_TECHNOLOGY_INPUT = "Much is made of what the likes of Facebook, Google and Apple know about users." + " Truth is, Amazon may know more. This technology observation is extracted from the internet."; protected static final String SUPER_STRONG_TECHNOLOGY_INPUT = "More than 400 million people trust Google with their e-mail, and 50 million store files" + " in the cloud using the Dropbox service. People manage their bank accounts, pay bills, trade stocks and " + "generally transfer or store huge volumes of personal data online. traveling seeks raises some questions Republican presidential. "; protected static final BytesRef TECHNOLOGY_RESULT = new BytesRef("technology"); protected RandomIndexWriter indexWriter; protected Directory dir; protected FieldType ft; protected String textFieldName; protected String categoryFieldName; protected String booleanFieldName; @Override @Before public void setUp() throws Exception { super.setUp(); dir = newDirectory(); indexWriter = new RandomIndexWriter(random(), dir); textFieldName = "text"; categoryFieldName = "cat"; booleanFieldName = "bool"; ft = new FieldType(TextField.TYPE_STORED); ft.setStoreTermVectors(true); ft.setStoreTermVectorOffsets(true); ft.setStoreTermVectorPositions(true); } @Override @After public void tearDown() throws Exception { super.tearDown(); indexWriter.close(); dir.close(); } protected ClassificationResult<T> checkCorrectClassification(Classifier<T> classifier, String inputDoc, T expectedResult) throws Exception { ClassificationResult<T> classificationResult = classifier.assignClass(inputDoc); T assignedClass = classificationResult.getAssignedClass(); assertNotNull(assignedClass); assertEquals("got an assigned class of " + assignedClass, expectedResult instanceof BytesRef ? ((BytesRef) expectedResult).utf8ToString() : expectedResult, assignedClass instanceof BytesRef ? ((BytesRef) assignedClass).utf8ToString() : assignedClass); double score = classificationResult.getScore(); assertTrue("score should be between 0 and 1, got:" + score, score <= 1 && score >= 0); return classificationResult; } protected void checkOnlineClassification(Classifier<T> classifier, String inputDoc, T expectedResult, Analyzer analyzer, String textFieldName, String classFieldName) throws Exception { checkOnlineClassification(classifier, inputDoc, expectedResult, analyzer, textFieldName, classFieldName, null); } protected void checkOnlineClassification(Classifier<T> classifier, String inputDoc, T expectedResult, Analyzer analyzer, String textFieldName, String classFieldName, Query query) throws Exception { getSampleIndex(analyzer); ClassificationResult<T> classificationResult = classifier.assignClass(inputDoc); assertNotNull(classificationResult.getAssignedClass()); assertEquals("got an assigned class of " + classificationResult.getAssignedClass(), expectedResult, classificationResult.getAssignedClass()); double score = classificationResult.getScore(); assertTrue("score should be between 0 and 1, got: " + score, score <= 1 && score >= 0); updateSampleIndex(); ClassificationResult<T> secondClassificationResult = classifier.assignClass(inputDoc); assertEquals(classificationResult.getAssignedClass(), secondClassificationResult.getAssignedClass()); assertEquals(Double.valueOf(score), Double.valueOf(secondClassificationResult.getScore())); } protected LeafReader getSampleIndex(Analyzer analyzer) throws IOException { indexWriter.close(); indexWriter = new RandomIndexWriter(random(), dir, newIndexWriterConfig(analyzer).setOpenMode(IndexWriterConfig.OpenMode.CREATE)); indexWriter.commit(); String text; Document doc = new Document(); text = "The traveling press secretary for Mitt Romney lost his cool and cursed at reporters " + "who attempted to ask questions of the Republican presidential candidate in a public plaza near the Tomb of " + "the Unknown Soldier in Warsaw Tuesday."; doc.add(new Field(textFieldName, text, ft)); doc.add(new Field(categoryFieldName, "politics", ft)); doc.add(new Field(booleanFieldName, "true", ft)); indexWriter.addDocument(doc); doc = new Document(); text = "Mitt Romney seeks to assure Israel and Iran, as well as Jewish voters in the United" + " States, that he will be tougher against Iran's nuclear ambitions than President Barack Obama."; doc.add(new Field(textFieldName, text, ft)); doc.add(new Field(categoryFieldName, "politics", ft)); doc.add(new Field(booleanFieldName, "true", ft)); indexWriter.addDocument(doc); doc = new Document(); text = "And there's a threshold question that he has to answer for the American people and " + "that's whether he is prepared to be commander-in-chief,\" she continued. \"As we look to the past events, we " + "know that this raises some questions about his preparedness and we'll see how the rest of his trip goes.\""; doc.add(new Field(textFieldName, text, ft)); doc.add(new Field(categoryFieldName, "politics", ft)); doc.add(new Field(booleanFieldName, "true", ft)); indexWriter.addDocument(doc); doc = new Document(); text = "Still, when it comes to gun policy, many congressional Democrats have \"decided to " + "keep quiet and not go there,\" said Alan Lizotte, dean and professor at the State University of New York at " + "Albany's School of Criminal Justice."; doc.add(new Field(textFieldName, text, ft)); doc.add(new Field(categoryFieldName, "politics", ft)); doc.add(new Field(booleanFieldName, "true", ft)); indexWriter.addDocument(doc); doc = new Document(); text = "Standing amongst the thousands of people at the state Capitol, Jorstad, director of " + "technology at the University of Wisconsin-La Crosse, documented the historic moment and shared it with the " + "world through the Internet."; doc.add(new Field(textFieldName, text, ft)); doc.add(new Field(categoryFieldName, "technology", ft)); doc.add(new Field(booleanFieldName, "false", ft)); indexWriter.addDocument(doc); doc = new Document(); text = "So, about all those experts and analysts who've spent the past year or so saying " + "Facebook was going to make a phone. A new expert has stepped forward to say it's not going to happen."; doc.add(new Field(textFieldName, text, ft)); doc.add(new Field(categoryFieldName, "technology", ft)); doc.add(new Field(booleanFieldName, "false", ft)); indexWriter.addDocument(doc); doc = new Document(); text = "More than 400 million people trust Google with their e-mail, and 50 million store files" + " in the cloud using the Dropbox service. People manage their bank accounts, pay bills, trade stocks and " + "generally transfer or store huge volumes of personal data online."; doc.add(new Field(textFieldName, text, ft)); doc.add(new Field(categoryFieldName, "technology", ft)); doc.add(new Field(booleanFieldName, "false", ft)); indexWriter.addDocument(doc); doc = new Document(); text = "unlabeled doc"; doc.add(new Field(textFieldName, text, ft)); indexWriter.addDocument(doc); indexWriter.commit(); indexWriter.forceMerge(1); return getOnlyLeafReader(indexWriter.getReader()); } protected LeafReader getRandomIndex(Analyzer analyzer, int size) throws IOException { indexWriter.close(); indexWriter = new RandomIndexWriter(random(), dir, newIndexWriterConfig(analyzer).setOpenMode(IndexWriterConfig.OpenMode.CREATE)); indexWriter.deleteAll(); indexWriter.commit(); FieldType ft = new FieldType(TextField.TYPE_STORED); ft.setStoreTermVectors(true); ft.setStoreTermVectorOffsets(true); ft.setStoreTermVectorPositions(true); Random random = random(); for (int i = 0; i < size; i++) { boolean b = random.nextBoolean(); Document doc = new Document(); doc.add(new Field(textFieldName, createRandomString(random), ft)); doc.add(new Field(categoryFieldName, String.valueOf(random.nextInt(1000)), ft)); doc.add(new Field(booleanFieldName, String.valueOf(b), ft)); indexWriter.addDocument(doc); } indexWriter.commit(); indexWriter.forceMerge(1); return getOnlyLeafReader(indexWriter.getReader()); } private String createRandomString(Random random) { StringBuilder builder = new StringBuilder(); for (int i = 0; i < 20; i++) { builder.append(TestUtil.randomSimpleString(random, 5)); builder.append(" "); } return builder.toString(); } private void updateSampleIndex() throws Exception { String text; Document doc = new Document(); text = "Warren Bennis says John F. Kennedy grasped a key lesson about the presidency that few have followed."; doc.add(new Field(textFieldName, text, ft)); doc.add(new Field(categoryFieldName, "politics", ft)); doc.add(new Field(booleanFieldName, "true", ft)); indexWriter.addDocument(doc); doc = new Document(); text = "Julian Zelizer says Bill Clinton is still trying to shape his party, years after the White House, while George W. Bush opts for a much more passive role."; doc.add(new Field(textFieldName, text, ft)); doc.add(new Field(categoryFieldName, "politics", ft)); doc.add(new Field(booleanFieldName, "true", ft)); indexWriter.addDocument(doc); doc = new Document(); text = "Crossfire: Sen. Tim Scott passes on Sen. Lindsey Graham endorsement"; doc.add(new Field(textFieldName, text, ft)); doc.add(new Field(categoryFieldName, "politics", ft)); doc.add(new Field(booleanFieldName, "true", ft)); indexWriter.addDocument(doc); doc = new Document(); text = "Illinois becomes 16th state to allow same-sex marriage."; doc.add(new Field(textFieldName, text, ft)); doc.add(new Field(categoryFieldName, "politics", ft)); doc.add(new Field(booleanFieldName, "true", ft)); indexWriter.addDocument(doc); doc = new Document(); text = "Apple is developing iPhones with curved-glass screens and enhanced sensors that detect different levels of pressure, according to a new report."; doc.add(new Field(textFieldName, text, ft)); doc.add(new Field(categoryFieldName, "technology", ft)); doc.add(new Field(booleanFieldName, "false", ft)); indexWriter.addDocument(doc); doc = new Document(); text = "The Xbox One is Microsoft's first new gaming console in eight years. It's a quality piece of hardware but it's also noteworthy because Microsoft is using it to make a statement."; doc.add(new Field(textFieldName, text, ft)); doc.add(new Field(categoryFieldName, "technology", ft)); doc.add(new Field(booleanFieldName, "false", ft)); indexWriter.addDocument(doc); doc = new Document(); text = "Google says it will replace a Google Maps image after a California father complained it shows the body of his teen-age son, who was shot to death in 2009."; doc.add(new Field(textFieldName, text, ft)); doc.add(new Field(categoryFieldName, "technology", ft)); doc.add(new Field(booleanFieldName, "false", ft)); indexWriter.addDocument(doc); doc = new Document(); text = "second unlabeled doc"; doc.add(new Field(textFieldName, text, ft)); indexWriter.addDocument(doc); indexWriter.commit(); } }