| /* |
| * 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); |
| assertNotNull(classificationResult); |
| 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(); |
| } |
| } |