package com.bbn.kbp.events2014.io; import com.bbn.bue.common.symbols.Symbol; import com.bbn.kbp.events2014.ArgumentOutput; import com.bbn.kbp.events2014.CorpusEventFrame; import com.bbn.kbp.events2014.CorpusEventLinking; import com.bbn.kbp.events2014.DocumentSystemOutput; import com.bbn.kbp.events2014.DocumentSystemOutput2015; import com.bbn.kbp.events2014.ResponseLinking; import com.google.common.base.Charsets; import com.google.common.base.Optional; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Sets; import com.google.common.io.Files; import java.io.File; import java.io.IOException; import java.util.Set; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Preconditions.checkState; public class SystemOutputStore2016 implements SystemOutputStore, CrossDocSystemOutputStore { private final Symbol systemID; private final ArgumentStore argumentStore; private final LinkingStore linkingStore; private final File corpusLinkingFile; private final CorpusEventFrameWriter eventFrameWriter = CorpusEventFrameIO.writerFor2016(); private final CorpusEventFrameLoader eventFrameReader = CorpusEventFrameIO.loaderFor2016(); private SystemOutputStore2016(final Symbol systemID, final ArgumentStore argStore, final LinkingStore linkingStore, File corpusLinkingFile) { this.systemID = checkNotNull(systemID); this.argumentStore = checkNotNull(argStore); this.linkingStore = checkNotNull(linkingStore); this.corpusLinkingFile = checkNotNull(corpusLinkingFile); } public static SystemOutputStore2016 open(File dir) throws IOException { return open(dir, AssessmentSpecFormats.Format.KBP2015); } static SystemOutputStore2016 open(File dir, final AssessmentSpecFormats.Format format) throws IOException { final Symbol systemID = Symbol.from(dir.getName()); final File argumentsDir = new File(dir, "arguments"); final File linkingDir = new File(dir, "linking"); final File corpusLinkingFile = new File(new File(dir, "corpusLinking"), "corpusLinking"); corpusLinkingFile.getParentFile().mkdirs(); final ArgumentStore argStore = AssessmentSpecFormats.openSystemOutputStore(argumentsDir, format); final LinkingStore linkingStore = LinkingStoreSource.createFor2016().openLinkingStore(linkingDir); if (argStore.docIDs().equals(linkingStore.docIDs())) { return new SystemOutputStore2016(systemID, argStore, linkingStore, corpusLinkingFile); } else { throw new RuntimeException("Argument and linking store docIDs do not match, missing " + Sets .symmetricDifference(argStore.docIDs(), linkingStore.docIDs())); } } public static SystemOutputStore2016 openOrCreate(final File dir) throws IOException { return openOrCreate(dir, AssessmentSpecFormats.Format.KBP2015); } static SystemOutputStore2016 openOrCreate(File dir, final AssessmentSpecFormats.Format format) throws IOException { final Symbol systemID = Symbol.from(dir.getName()); final File argumentsDir = new File(dir, "arguments"); final File linkingDir = new File(dir, "linking"); final File corpusLinkingFile = new File(new File(dir, "corpusLinking"), "corpusLinking"); corpusLinkingFile.getParentFile().mkdirs(); final ArgumentStore argStore = AssessmentSpecFormats.openOrCreateSystemOutputStore(argumentsDir, format); final LinkingStore linkingStore = LinkingStoreSource.createFor2016().openOrCreateLinkingStore(linkingDir); if (argStore.docIDs().equals(linkingStore.docIDs())) { return new SystemOutputStore2016(systemID, argStore, linkingStore, corpusLinkingFile); } else { throw new RuntimeException("Argument and linking store docIDs do not match, missing " + Sets .symmetricDifference(argStore.docIDs(), linkingStore.docIDs())); } } @Override public CorpusEventLinking readCorpusEventFrames() throws IOException { if (corpusLinkingFile.isFile()) { return eventFrameReader.loadCorpusEventFrames(Files.asCharSource(corpusLinkingFile, Charsets.UTF_8)); } else { return CorpusEventLinking.of(ImmutableSet.<CorpusEventFrame>of()); } } @Override public void writeCorpusEventFrames(CorpusEventLinking corpusEventFrames) throws IOException { corpusLinkingFile.getParentFile().mkdirs(); eventFrameWriter.writeCorpusEventFrames(corpusEventFrames, Files.asCharSink(corpusLinkingFile, Charsets.UTF_8)); } @Override public final Symbol systemID() { return systemID; } @Override public final Set<Symbol> docIDs() throws IOException { return argumentStore.docIDs(); } @Override public DocumentSystemOutput2015 read(final Symbol docID) throws IOException { final ArgumentOutput argOutput = argumentStore.read(docID); final Optional<ResponseLinking> linking = linkingStore.read(argOutput); checkState(linking.isPresent(), "No linking found for " + docID); return DocumentSystemOutput2015.from(argOutput, linking.get()); } @Override public void write(final DocumentSystemOutput output) throws IOException { checkArgument(output instanceof DocumentSystemOutput2015); argumentStore.write(output.arguments()); linkingStore.write(((DocumentSystemOutput2015) output).linking()); } @Override public void close() throws IOException { argumentStore.close(); linkingStore.close(); } }