wrap base config in PaLMConfig (#15)

Author: dzlab

llm4j-palm/src/main/java/org/llm4j/palm/PaLMConfig.java

@@ -1,5 +1,51 @@
 package org.llm4j.palm;
 
+import org.apache.commons.configuration2.Configuration;
+
 public class PaLMConfig {
     public static String MODEL_ID_KEY = "palm.modelId";
+    public static String CHAT_MODEL_ID = "models/chat-bison-001";
+    public static String EMBED_MODEL_ID = "models/embedding-gecko-001";
+    public static String TEXT_MODEL_ID = "models/text-bison-001";
+
+    private final Configuration configs;
+    public PaLMConfig(Configuration configs) {
+        this.configs = configs;
+    }
+
+    public String getChatModelId() {
+        return configs.getString(PaLMConfig.MODEL_ID_KEY, CHAT_MODEL_ID);
+    }
+
+    public String getEmbedModelId() {
+        return configs.getString(PaLMConfig.MODEL_ID_KEY, EMBED_MODEL_ID);
+    }
+
+    public String getTextModelId() {
+        return configs.getString(PaLMConfig.MODEL_ID_KEY, TEXT_MODEL_ID);
+    }
+
+    public Integer getTopK() {
+        return configs.getInteger("topK", 1);
+    }
+
+    public Double getTopP() {
+        return configs.getDouble("topP", 0.9);
+    }
+
+    public Double getTemperature() {
+        return configs.getDouble("temperature", 1.0);
+    }
+
+    public Integer getMaxNewTokens() {
+        return configs.getInteger("maxNewTokens", 10);
+    }
+
+    public Integer getMaxOutputTokens() {
+        return configs.getInteger("maxOutputTokens", 1000);
+    }
+
+    public Integer getCandidateCount() {
+        return configs.getInteger("candidateCount", 1);
+    }
 }

llm4j-palm/src/main/java/org/llm4j/palm/PaLMLanguageModel.java

@@ -14,7 +14,7 @@
 
 public class PaLMLanguageModel implements LanguageModel {
 
-    private Configuration config;
+    private PaLMConfig config;
     private PaLMClient client;
 
     public PaLMLanguageModel(Builder builder) {
@@ -64,11 +64,11 @@ public List<Float> embed(String text) {
     }
 
     public static final class Builder implements LanguageModelFactory {
-        private Configuration config;
+        private PaLMConfig config;
         private PaLMClient client;
 
         public LanguageModel getLanguageModel(Configuration config) {
-            this.config = config;
+            this.config = new PaLMConfig(config);
             this.client = new PaLMClient.Builder().withConfig(config).build();
             return new PaLMLanguageModel(this);
         }

llm4j-palm/src/main/java/org/llm4j/palm/request/EmbedRequestFactory.java

@@ -5,7 +5,6 @@
 import org.llm4j.palm.PaLMConfig;
 
 public class EmbedRequestFactory {
-    static String EMBED_MODEL_ID = "models/embedding-gecko-001";
 
     /**
      * which model to use to generate the result
@@ -19,8 +18,8 @@ public EmbedRequestFactory withText(String text) {
         return this;
     }
 
-    public EmbedRequestFactory withConfig(Configuration configs) {
-        this.modelId = configs.getString(PaLMConfig.MODEL_ID_KEY, EMBED_MODEL_ID);
+    public EmbedRequestFactory withConfig(PaLMConfig configs) {
+        this.modelId = configs.getEmbedModelId();
         return this;
     }
 

llm4j-palm/src/main/java/org/llm4j/palm/request/MessageRequestFactory.java

@@ -9,8 +9,6 @@
 
 public class MessageRequestFactory {
 
-    static String CHAT_MODEL_ID = "models/chat-bison-001";
-
     /**
      * which model to use to generate the result
      */
@@ -56,8 +54,8 @@ public MessageRequestFactory withContext(String context) {
         return this;
     }
 
-    public MessageRequestFactory withConfig(Configuration configs) {
-        this.modelId = configs.getString(PaLMConfig.MODEL_ID_KEY, CHAT_MODEL_ID);
+    public MessageRequestFactory withConfig(PaLMConfig configs) {
+        this.modelId = configs.getChatModelId();
         this.parameters = PaLMRequestParameters.builder()
                         .withConfig(configs)
                         .build();

llm4j-palm/src/main/java/org/llm4j/palm/request/PaLMRequestParameters.java

@@ -1,6 +1,7 @@
 package org.llm4j.palm.request;
 
 import org.apache.commons.configuration2.Configuration;
+import org.llm4j.palm.PaLMConfig;
 
 public class PaLMRequestParameters {
 
@@ -42,13 +43,13 @@ public static final class Builder {
         private Integer maxOutputTokens;
         private Integer candidateCount;
 
-        public Builder withConfig(Configuration configs) {
-            this.topK = configs.getInteger("topK", 1);
-            this.topP = configs.getDouble("topP", 0.9);
-            this.temperature = configs.getDouble("temperature", 1.0);
-            this.maxNewTokens = configs.getInteger("maxNewTokens", 10);
-            this.maxOutputTokens = configs.getInteger("maxOutputTokens", 1000);
-            this.candidateCount = configs.getInteger("candidateCount", 1);
+        public Builder withConfig(PaLMConfig configs) {
+            this.topK = configs.getTopK();
+            this.topP = configs.getTopP();
+            this.temperature = configs.getTemperature();
+            this.maxNewTokens = configs.getMaxNewTokens();
+            this.maxOutputTokens = configs.getMaxOutputTokens();
+            this.candidateCount = configs.getCandidateCount();
             return this;
         }
         public PaLMRequestParameters build() {

llm4j-palm/src/main/java/org/llm4j/palm/request/TextRequestFactory.java

@@ -2,11 +2,9 @@
 
 import com.google.ai.generativelanguage.v1beta2.GenerateTextRequest;
 import com.google.ai.generativelanguage.v1beta2.TextPrompt;
-import org.apache.commons.configuration2.Configuration;
 import org.llm4j.palm.PaLMConfig;
 
 public class TextRequestFactory {
-    static String TEXT_MODEL_ID = "models/text-bison-001";
 
     /**
      * which model to use to generate the result
@@ -22,8 +20,8 @@ public TextRequestFactory withInputs(String inputs) {
         return this;
     }
 
-    public TextRequestFactory withConfig(Configuration configs) {
-        this.modelId = configs.getString(PaLMConfig.MODEL_ID_KEY, TEXT_MODEL_ID);
+    public TextRequestFactory withConfig(PaLMConfig configs) {
+        this.modelId = configs.getTextModelId();
         this.parameters = PaLMRequestParameters.builder()
                         .withConfig(configs)
                         .build();