Add support for the model option

mziccard · mziccard · commit a5962258d050 · 2016-11-11T17:01:38.000+01:00
diff --git a/google-cloud-translate/src/main/java/com/google/cloud/translate/Translate.java b/google-cloud-translate/src/main/java/com/google/cloud/translate/Translate.java
@@ -81,6 +81,19 @@ public static TranslateOption sourceLanguage(String sourceLanguage) {
     public static TranslateOption targetLanguage(String targetLanguage) {
       return new TranslateOption(TranslateRpc.Option.TARGET_LANGUAGE, targetLanguage);
     }
+
+    /**
+     * Sets the language translation model. You can use this parameter to take advantage of Neural
+     * Machine Translation. Possible values are {@code base} and {@code nmt}. Google Translate could
+     * use a different model to translate your text, use {@link Translation#getModel()} to know
+     * which model was used for translation. Please notice that you must be whitelisted to use this
+     * parameter.
+     *
+     * @param model the language translation model
+     */
+    public static TranslateOption model(String model) {
+      return new TranslateOption(TranslateRpc.Option.MODEL, model);
+    }
   }
 
   /**
diff --git a/google-cloud-translate/src/main/java/com/google/cloud/translate/Translation.java b/google-cloud-translate/src/main/java/com/google/cloud/translate/Translation.java
@@ -21,6 +21,7 @@
 import com.google.common.base.MoreObjects;
 
 import java.io.Serializable;
+import java.util.List;
 import java.util.Objects;
 
 /**
@@ -43,10 +44,12 @@ public Translation apply(TranslationsResource translationPb) {
 
   private final String translatedText;
   private final String sourceLanguage;
+  private final String model;
 
-  private Translation(String translatedText, String sourceLanguage) {
+  private Translation(String translatedText, String sourceLanguage, String model) {
     this.translatedText = translatedText;
     this.sourceLanguage = sourceLanguage;
+    this.model = model;
   }
 
   /**
@@ -81,6 +84,18 @@ public String getSourceLanguage() {
     return sourceLanguage;
   }
 
+  /**
+   * Returns the translation model used to translate the text. This value is only available if
+   * {@link Translate.TranslateOption#model(String)} was passed to
+   * {@link Translate#translate(List, Translate.TranslateOption...)}.
+   *
+   * <p>Please notice that {@link Translate.TranslateOption#model(String)} is only available to
+   * whitelisted users.
+   */
+  public String getModel() {
+    return model;
+  }
+
   @Override
   public String toString() {
     return MoreObjects.toStringHelper(this)
@@ -108,7 +123,8 @@ public final boolean equals(Object obj) {
   }
 
   static Translation fromPb(TranslationsResource translationPb) {
+    // todo remove get("model") as soon as REST apiary supports model
     return new Translation(translationPb.getTranslatedText(),
-        translationPb.getDetectedSourceLanguage());
+        translationPb.getDetectedSourceLanguage(), (String) translationPb.get("model"));
   }
 }
diff --git a/google-cloud-translate/src/main/java/com/google/cloud/translate/spi/DefaultTranslateRpc.java b/google-cloud-translate/src/main/java/com/google/cloud/translate/spi/DefaultTranslateRpc.java
@@ -16,6 +16,7 @@
 
 package com.google.cloud.translate.spi;
 
+import static com.google.cloud.translate.spi.TranslateRpc.Option.MODEL;
 import static com.google.cloud.translate.spi.TranslateRpc.Option.SOURCE_LANGUAGE;
 import static com.google.cloud.translate.spi.TranslateRpc.Option.TARGET_LANGUAGE;
 import static com.google.common.base.MoreObjects.firstNonNull;
@@ -118,13 +119,17 @@ public List<LanguagesResource> listSupportedLanguages(Map<Option, ?> optionMap)
   public List<TranslationsResource> translate(List<String> texts, Map<Option, ?> optionMap) {
     try {
       final String sourceLanguage = SOURCE_LANGUAGE.getString(optionMap);
+      String model = MODEL.getString(optionMap);
       ImmutableMap.Builder<String, Object> contentBuilder = ImmutableMap.builder();
       contentBuilder.put("target",
           firstNonNull(TARGET_LANGUAGE.getString(optionMap), options.getTargetLanguage()));
       contentBuilder.put("q", texts);
       if (sourceLanguage != null) {
         contentBuilder.put("source", sourceLanguage);
       }
+      if (model != null) {
+        contentBuilder.put("model", model);
+      }
       HttpRequest httpRequest = translate.getRequestFactory()
           .buildPostRequest(buildTargetUrl(""),
               new JsonHttpContent(translate.getJsonFactory(), contentBuilder.build()))
diff --git a/google-cloud-translate/src/main/java/com/google/cloud/translate/spi/TranslateRpc.java b/google-cloud-translate/src/main/java/com/google/cloud/translate/spi/TranslateRpc.java
@@ -27,7 +27,8 @@ public interface TranslateRpc {
 
   enum Option {
     SOURCE_LANGUAGE("source"),
-    TARGET_LANGUAGE("target");
+    TARGET_LANGUAGE("target"),
+    MODEL("model");
 
     private final String value;
 
diff --git a/google-cloud-translate/src/test/java/com/google/cloud/translate/TranslateImplTest.java b/google-cloud-translate/src/test/java/com/google/cloud/translate/TranslateImplTest.java
@@ -86,9 +86,11 @@ public class TranslateImplTest {
       TranslateOption.targetLanguage("en");
   private static final TranslateOption SOURCE_LANGUAGE_OPTION =
       TranslateOption.sourceLanguage("de");
+  private static final TranslateOption MODEL_OPTION = TranslateOption.model("nmt");
   private static final Map<TranslateRpc.Option, ?> TRANSLATE_OPTIONS = ImmutableMap.of(
       TranslateRpc.Option.TARGET_LANGUAGE, TARGET_LANGUAGE_OPTION.getValue(),
-      TranslateRpc.Option.SOURCE_LANGUAGE, SOURCE_LANGUAGE_OPTION.getValue());
+      TranslateRpc.Option.SOURCE_LANGUAGE, SOURCE_LANGUAGE_OPTION.getValue(),
+      TranslateRpc.Option.MODEL, "nmt");
 
   private TranslateOptions options;
   private TranslateRpcFactory rpcFactoryMock;
@@ -285,7 +287,7 @@ public void testTranslateWithOptions() {
     EasyMock.replay(translateRpcMock);
     initializeService();
     assertEquals(TRANSLATION2,
-        translate.translate(text, TARGET_LANGUAGE_OPTION, SOURCE_LANGUAGE_OPTION));
+        translate.translate(text, TARGET_LANGUAGE_OPTION, SOURCE_LANGUAGE_OPTION, MODEL_OPTION));
   }
 
   @Test
@@ -309,7 +311,7 @@ public void testTranslateListWithOptions() {
     EasyMock.replay(translateRpcMock);
     initializeService();
     assertEquals(ImmutableList.of(TRANSLATION2),
-        translate.translate(texts, TARGET_LANGUAGE_OPTION, SOURCE_LANGUAGE_OPTION));
+        translate.translate(texts, TARGET_LANGUAGE_OPTION, SOURCE_LANGUAGE_OPTION, MODEL_OPTION));
   }
 
   @Test
