Add the ability to read & write a fast coref model

Author: AngledLuffa

src/edu/stanford/nlp/coref/fastneural/FastNeuralCorefModel.java

@@ -2,6 +2,7 @@
 
 import java.io.Serializable;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -37,8 +38,9 @@ public class FastNeuralCorefModel implements Serializable {
   private List<SimpleMatrix> networkLayers;
 
   public FastNeuralCorefModel(EmbeddingExtractor embeddingExtractor,
-      Map<String, Integer> pairFeatureIds, Map<String, Integer> mentionFeatureIds,
-      List<SimpleMatrix> weights) {
+                              Map<String, Integer> pairFeatureIds,
+                              Map<String, Integer> mentionFeatureIds,
+                              List<SimpleMatrix> weights) {
     this.embeddingExtractor = embeddingExtractor;
     this.pairFeatureIds = pairFeatureIds;
     this.mentionFeatureIds = mentionFeatureIds;
@@ -53,6 +55,31 @@ public FastNeuralCorefModel(EmbeddingExtractor embeddingExtractor,
     networkLayers = new ArrayList<>(weights.subList(7, weights.size()));
   }
 
+  public EmbeddingExtractor getEmbeddingExtractor() {
+    return embeddingExtractor;
+  }
+
+  public Map<String, Integer> getPairFeatureIds() {
+    return Collections.unmodifiableMap(pairFeatureIds);
+  }
+
+  public Map<String, Integer> getMentionFeatureIds() {
+    return Collections.unmodifiableMap(mentionFeatureIds);
+  }
+
+  public List<SimpleMatrix> getAllWeights() {
+    List<SimpleMatrix> weights = new ArrayList<>();
+    weights.add(anaphorKernel);
+    weights.add(anaphorBias);
+    weights.add(antecedentKernel);
+    weights.add(anaphorBias);
+    weights.add(pairFeaturesKernel);
+    weights.add(pairFeaturesBias);
+    weights.add(NARepresentation);
+    weights.addAll(networkLayers);
+    return Collections.unmodifiableList(weights);
+  }
+
   public double score(Mention antecedent, Mention anaphor, Counter<String> antecedentFeatures,
       Counter<String> anaphorFeatures, Counter<String> pairFeatures,
       Map<Integer, SimpleMatrix> antecedentCache, Map<Integer, SimpleMatrix> anaphorCache) {

src/edu/stanford/nlp/coref/neural/EmbeddingExtractor.java

@@ -28,13 +28,29 @@ public class EmbeddingExtractor implements Serializable {
   private final String naEmbedding;
 
   public EmbeddingExtractor(boolean conll, Embedding staticWordEmbeddings,
-      Embedding tunedWordEmbeddings, String naEmbedding) {
+                            Embedding tunedWordEmbeddings, String naEmbedding) {
     this.conll = conll;
     this.staticWordEmbeddings = staticWordEmbeddings;
     this.tunedWordEmbeddings = tunedWordEmbeddings;
     this.naEmbedding = naEmbedding;
   }
 
+  public boolean isConll() {
+    return conll;
+  }
+
+  public Embedding getStaticWordEmbeddings() {
+    return staticWordEmbeddings;
+  }
+
+  public Embedding getTunedWordEmbeddings() {
+    return tunedWordEmbeddings;
+  }
+
+  public String getNAEmbedding() {
+    return naEmbedding;
+  }
+
   public SimpleMatrix getDocumentEmbedding(Document document) {
     if (!conll) {
       return new SimpleMatrix(staticWordEmbeddings.getEmbeddingSize(), 1);

src/edu/stanford/nlp/neural/ConvertModels.java

@@ -14,6 +14,8 @@
 
 import org.ejml.simple.SimpleMatrix;
 
+import edu.stanford.nlp.coref.fastneural.FastNeuralCorefModel;
+import edu.stanford.nlp.coref.neural.EmbeddingExtractor;
 import edu.stanford.nlp.coref.neural.NeuralCorefModel;
 import edu.stanford.nlp.io.IOUtils;
 import edu.stanford.nlp.parser.dvparser.DVModel;
@@ -34,7 +36,7 @@ public enum Stage {
   }
 
   public enum Model {
-    SENTIMENT, DVPARSER, COREF, EMBEDDING
+    SENTIMENT, DVPARSER, COREF, EMBEDDING, FASTCOREF
   }
 
   /**
@@ -222,6 +224,7 @@ public static void writeEmbedding(Embedding embedding, ObjectOutputStream out)
   public static Embedding readEmbedding(ObjectInputStream in)
     throws IOException, ClassNotFoundException
   {
+
     Function<List<List<Double>>, SimpleMatrix> f = (x) -> toMatrix(x);
     Map<String, List<List<Double>>> map = ErasureUtils.uncheckedCast(in.readObject());
     Map<String, SimpleMatrix> vectors = transformMap(map, f);
@@ -267,6 +270,48 @@ public static NeuralCorefModel readCoref(ObjectInputStream in)
     return model;
   }
 
+  public static void writeFastCoref(FastNeuralCorefModel model, ObjectOutputStream out)
+    throws IOException
+  {
+    Function<SimpleMatrix, List<List<Double>>> f = (SimpleMatrix x) -> fromMatrix(x);
+
+    EmbeddingExtractor embedding = model.getEmbeddingExtractor();
+    out.writeObject(embedding.isConll());
+    Embedding staticEmbedding = embedding.getStaticWordEmbeddings();
+    if (staticEmbedding == null) {
+      out.writeObject(false);
+    } else {
+      out.writeObject(true);
+      writeEmbedding(staticEmbedding, out);
+    }
+    writeEmbedding(embedding.getTunedWordEmbeddings(), out);
+    out.writeObject(embedding.getNAEmbedding());
+
+    out.writeObject(model.getPairFeatureIds());
+    out.writeObject(model.getMentionFeatureIds());
+    out.writeObject(CollectionUtils.transformAsList(model.getAllWeights(), f));
+  }
+
+  public static FastNeuralCorefModel readFastCoref(ObjectInputStream in)
+    throws IOException, ClassNotFoundException
+  {
+    Function<List<List<Double>>, SimpleMatrix> f = (x) -> toMatrix(x);
+
+    boolean conll = ErasureUtils.uncheckedCast(in.readObject());
+    boolean hasStatic = ErasureUtils.uncheckedCast(in.readObject());
+    Embedding staticEmbedding = (hasStatic) ? readEmbedding(in) : null;
+    Embedding tunedEmbedding = readEmbedding(in);
+    String naEmbedding = ErasureUtils.uncheckedCast(in.readObject());
+
+    EmbeddingExtractor embedding = new EmbeddingExtractor(conll, staticEmbedding, tunedEmbedding, naEmbedding);
+
+    Map<String, Integer> pairFeatures = ErasureUtils.uncheckedCast(in.readObject());
+    Map<String, Integer> mentionFeatures = ErasureUtils.uncheckedCast(in.readObject());
+    List<SimpleMatrix> weights = CollectionUtils.transformAsList(ErasureUtils.uncheckedCast(in.readObject()), f);
+
+    return new FastNeuralCorefModel(embedding, pairFeatures, mentionFeatures, weights);
+  }
+
   /**
    * This program converts a sentiment model or an RNN parser model
    * from EJML v23, used by CoreNLP 3.9.2, to a more recent version of
@@ -314,6 +359,12 @@ public static NeuralCorefModel readCoref(ObjectInputStream in)
    * <br>
    * <code> java edu.stanford.nlp.neural.ConvertModels -stage NEW -model EMBEDDING -input /scr/nlp/data/coref/models/neural/english/english-embeddings.INT.ser.gz -output /scr/nlp/data/coref/models/neural/english/english-embeddings.e39.ser.gz</code>
    * <br>
+   * There is another coref model which isn't used in corenlp, but it might be in the future.  To upgrade this, use <code>-model FASTCOREF</code>
+   * <br>
+   * <code> java edu.stanford.nlp.neural.ConvertModels -stage OLD -model FASTCOREF -input /scr/nlp/data/coref/models/fastneural/fast-english-model.e38.ser.gz -output /scr/nlp/data/coref/models/fastneural/fast-english-model.INT.ser.gz</code>
+   * <br>
+   * <code> java edu.stanford.nlp.neural.ConvertModels -stage NEW -model FASTCOREF -input /scr/nlp/data/coref/models/fastneural/fast-english-model.INT.ser.gz -output /scr/nlp/data/coref/models/fastneural/fast-english-model.e39.ser.gz</code>
+   * <br>
    *
    * @author <a href=horatio@gmail.com>John Bauer</a>
    */
@@ -331,7 +382,7 @@ public static void main(String[] args) throws IOException, ClassNotFoundExceptio
     try {
       modelType = Model.valueOf(props.getProperty("model").toUpperCase());
     } catch (IllegalArgumentException | NullPointerException e) {
-      throw new IllegalArgumentException("Please specify -model, either SENTIMENT, DVPARSER, EMBEDDING, COREF");
+      throw new IllegalArgumentException("Please specify -model, either SENTIMENT, DVPARSER, EMBEDDING, COREF, FASTCOREF");
     }
 
     if (!props.containsKey("input")) {
@@ -399,6 +450,18 @@ public static void main(String[] args) throws IOException, ClassNotFoundExceptio
         in.close();
         IOUtils.writeObjectToFile(model, outputPath);
       }
+    } else if (modelType == Model.FASTCOREF) {
+      if (stage == Stage.OLD) {
+        FastNeuralCorefModel model = ErasureUtils.uncheckedCast(IOUtils.readObjectFromURLOrClasspathOrFileSystem(inputPath));
+        ObjectOutputStream out = IOUtils.writeStreamFromString(outputPath);
+        writeFastCoref(model, out);
+        out.close();
+      } else {
+        ObjectInputStream in = IOUtils.readStreamFromString(inputPath);
+        FastNeuralCorefModel model = readFastCoref(in);
+        in.close();
+        IOUtils.writeObjectToFile(model, outputPath);
+      }
     } else {
       throw new IllegalArgumentException("Unknown model type " + modelType);
     }