Move parameters to GenerationConfig

google-gemini · Sep 26, 2024 · 3c959e9 · 3c959e9
1 parent 1e22275
commit 3c959e9
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 63 deletions.
diff --git a/src/methods/generate-content.test.ts b/src/methods/generate-content.test.ts
@@ -42,10 +42,6 @@ const fakeRequestParams: GenerateContentRequest = {
       threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
     },
   ],
-  presencePenalty: 0.5,
-  frequencyPenalty: 0.1,
-  responseLogProbs: true,
-  logProbs: 2,
 };
 
 describe("generateContent()", () => {

diff --git a/src/models/generative-model.test.ts b/src/models/generative-model.test.ts
@@ -71,10 +71,6 @@ describe("GenerativeModel", () => {
             threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
           },
         ],
-        presencePenalty: 0.6,
-        frequencyPenalty: 0.5,
-        responseLogProbs: true,
-        logProbs: 2,
         tools: [{ functionDeclarations: [{ name: "myfunc" }] }],
         toolConfig: {
           functionCallingConfig: { mode: FunctionCallingMode.NONE },
@@ -96,9 +92,6 @@ describe("GenerativeModel", () => {
       genModel.generationConfig?.responseSchema.properties.testField.type,
     ).to.equal(SchemaType.STRING);
     expect(genModel.safetySettings?.length).to.equal(1);
-    expect(genModel.presencePenalty).to.equal(0.6);
-    expect(genModel.frequencyPenalty).to.equal(0.5);
-    expect(genModel.tools?.length).to.equal(1);
     expect(genModel.toolConfig?.functionCallingConfig.mode).to.equal(
       FunctionCallingMode.NONE,
     );
@@ -122,11 +115,7 @@ describe("GenerativeModel", () => {
           value.includes("be friendly") &&
           value.includes("temperature") &&
           value.includes("testField") &&
-          value.includes(HarmBlockThreshold.BLOCK_LOW_AND_ABOVE) &&
-          value.includes("presencePenalty") &&
-          value.includes("frequencyPenalty") &&
-          value.includes("responseLogProbs") &&
-          value.includes("logProbs")
+          value.includes(HarmBlockThreshold.BLOCK_LOW_AND_ABOVE)
         );
       }),
       match((value) => {
@@ -220,10 +209,6 @@ describe("GenerativeModel", () => {
           threshold: HarmBlockThreshold.BLOCK_NONE,
         },
       ],
-      presencePenalty: 0.6,
-      frequencyPenalty: 0.5,
-      responseLogProbs: true,
-      logProbs: 2,
       contents: [{ role: "user", parts: [{ text: "hello" }] }],
       tools: [{ functionDeclarations: [{ name: "otherfunc" }] }],
       toolConfig: { functionCallingConfig: { mode: FunctionCallingMode.AUTO } },
@@ -243,10 +228,6 @@ describe("GenerativeModel", () => {
           value.includes("newTestField") &&
           !value.includes("testField") &&
           value.includes(HarmCategory.HARM_CATEGORY_HARASSMENT) &&
-          value.includes("presencePenalty") &&
-          value.includes("frequencyPenalty") &&
-          value.includes("responseLogProbs") &&
-          value.includes("logProbs")
         );
       }),
       {},

diff --git a/src/models/generative-model.ts b/src/models/generative-model.ts
@@ -63,10 +63,6 @@ export class GenerativeModel {
   toolConfig?: ToolConfig;
   systemInstruction?: Content;
   cachedContent: CachedContent;
-  presencePenalty?: number;
-  frequencyPenalty?: number;
-  responseLogProbs?: boolean;
-  logProbs?: number;
 
   constructor(
     public apiKey: string,
@@ -88,10 +84,6 @@ export class GenerativeModel {
       modelParams.systemInstruction,
     );
     this.cachedContent = modelParams.cachedContent;
-    this.presencePenalty = modelParams.presencePenalty;
-    this.frequencyPenalty = modelParams.frequencyPenalty;
-    this.responseLogProbs = modelParams.responseLogProbs;
-    this.logProbs = modelParams.logProbs;
   }
 
   /**
@@ -121,10 +113,6 @@ export class GenerativeModel {
         toolConfig: this.toolConfig,
         systemInstruction: this.systemInstruction,
         cachedContent: this.cachedContent?.name,
-        presencePenalty: this.presencePenalty,
-        frequencyPenalty: this.frequencyPenalty,
-        responseLogProbs: this.responseLogProbs,
-        logProbs: this.logProbs,
         ...formattedParams,
       },
       generativeModelRequestOptions,
@@ -160,10 +148,6 @@ export class GenerativeModel {
         toolConfig: this.toolConfig,
         systemInstruction: this.systemInstruction,
         cachedContent: this.cachedContent?.name,
-        presencePenalty: this.presencePenalty,
-        frequencyPenalty: this.frequencyPenalty,
-        responseLogProbs: this.responseLogProbs,
-        logProbs: this.logProbs,
         ...formattedParams,
       },
       generativeModelRequestOptions,
@@ -210,10 +194,6 @@ export class GenerativeModel {
       toolConfig: this.toolConfig,
       systemInstruction: this.systemInstruction,
       cachedContent: this.cachedContent,
-      presencePenalty: this.presencePenalty,
-      frequencyPenalty: this.frequencyPenalty,
-      responseLogProbs: this.responseLogProbs,
-      logProbs: this.logProbs,
     });
     const generativeModelRequestOptions: SingleRequestOptions = {
       ...this._requestOptions,

diff --git a/types/requests.ts b/types/requests.ts
@@ -31,25 +31,6 @@ import {
 export interface BaseParams {
   safetySettings?: SafetySetting[];
   generationConfig?: GenerationConfig;
-  /**
-   * Presence penalty applied to the next token's logprobs if the token has
-   * already been seen in the response.
-   */
-  presencePenalty?: number;
-  /**
-   * Frequency penalty applied to the next token's logprobs, multiplied by the
-   * number of times each token has been seen in the respponse so far.
-   */
-  frequencyPenalty?: number;
-  /**
-   * If True, export the logprobs results in response.
-   */
-  responseLogProbs?: boolean;
-  /**
-   * Valid if responseLogProbs is set to True. This will set the number of top
-   * logprobs to return at each decoding step in the logProbsResult.
-   */
-  logProbs?: number;
 }
 
 /**
@@ -121,6 +102,25 @@ export interface GenerationConfig {
    * this is limited to `application/json`.
    */
   responseSchema?: ResponseSchema;
+  /**
+   * Presence penalty applied to the next token's logprobs if the token has
+   * already been seen in the response.
+   */
+  presencePenalty?: number;
+  /**
+   * Frequency penalty applied to the next token's logprobs, multiplied by the
+   * number of times each token has been seen in the respponse so far.
+   */
+  frequencyPenalty?: number;
+  /**
+   * If True, export the logprobs results in response.
+   */
+  responseLogProbs?: boolean;
+  /**
+   * Valid if responseLogProbs is set to True. This will set the number of top
+   * logprobs to return at each decoding step in the logProbsResult.
+   */
+  logProbs?: number;
 }
 
 /**