Sampling: Add forced and additive support

Forced overrides overwrite what's passed in to that specific value.
Additive adds the override to an existing array.

Signed-off-by: kingbri <8082010+kingbri1@users.noreply.github.com>

Author: kingbri1

common/samplerOverrides.ts

@@ -1,6 +1,7 @@
 import * as YAML from "@std/yaml";
 import * as z from "@/common/myZod.ts";
 import { logger } from "@/common/logging.ts";
+import { BaseSamplerRequest } from "@/common/sampling.ts";
 
 export const SamplerOverride = z.object({
     override: z.unknown().refine((val) => val !== undefined && val !== null, {
@@ -16,17 +17,28 @@ export type SamplerOverride = z.infer<typeof SamplerOverride>;
 class SamplerOverridesContainer {
     selectedPreset?: string;
     overrides: Record<string, SamplerOverride> = {};
+    forcedOverrides: Record<string, SamplerOverride> = {};
 }
 
-export const overridesContainer = new SamplerOverridesContainer();
+// No need to export this, the functions work properly
+const overridesContainer = new SamplerOverridesContainer();
 
 export function overridesFromDict(newOverrides: Record<string, unknown>) {
     const parsedOverrides: Record<string, SamplerOverride> = {};
 
+    // Forced also includes additive
+    const forcedOverrides: Record<string, SamplerOverride> = {};
+
     // Validate each entry as a SamplerOverride type
     for (const [key, value] of Object.entries(newOverrides)) {
         try {
-            parsedOverrides[key] = SamplerOverride.parse(value);
+            const parsedOverride = SamplerOverride.parse(value);
+            parsedOverrides[key] = parsedOverride;
+
+            // Add to forced object for faster lookup
+            if (parsedOverride.force || parsedOverride.additive) {
+                forcedOverrides[key] = parsedOverride;
+            }
         } catch (error) {
             if (error instanceof Error) {
                 logger.error(error.stack);
@@ -39,6 +51,7 @@ export function overridesFromDict(newOverrides: Record<string, unknown>) {
     }
 
     overridesContainer.overrides = parsedOverrides;
+    overridesContainer.forcedOverrides = forcedOverrides;
 }
 
 export async function overridesFromFile(presetName: string) {
@@ -62,6 +75,27 @@ export async function overridesFromFile(presetName: string) {
     }
 }
 
+export function forcedSamplerOverrides(params: BaseSamplerRequest) {
+    const castParams = params as Record<string, unknown>;
+
+    for (
+        const [key, value] of Object.entries(overridesContainer.forcedOverrides)
+    ) {
+        if (value.force) {
+            castParams[key] = value.override;
+        } else if (
+            value.additive && Array.isArray(value.override) &&
+            Array.isArray(castParams[key])
+        ) {
+            castParams[key] = Array.from(
+                new Set([...castParams[key], ...value.override]),
+            );
+        }
+    }
+
+    return params;
+}
+
 export function getSamplerDefault<T>(key: string, fallback?: T): T {
     const defaultValue = overridesContainer.overrides[key]?.override ??
         fallback;

common/sampling.ts

@@ -1,4 +1,5 @@
 import * as z from "@/common/myZod.ts";
+import { forcedSamplerOverrides } from "@/common/samplerOverrides.ts";
 
 // Sampling schemas
 const GenerationOptionsSchema = z.aliasedObject(
@@ -217,8 +218,8 @@ const MirostatSchema = z.object({
         description: "Mirostat options",
     });
 
-// Construct from aliased sampler requests
-export const BaseSamplerRequest = GenerationOptionsSchema
+// Define the schema
+const BaseSamplerRequestSchema = GenerationOptionsSchema
     .and(TemperatureSamplerSchema)
     .and(AlphabetSamplerSchema)
     .and(PenaltySamplerSchema)
@@ -227,4 +228,11 @@ export const BaseSamplerRequest = GenerationOptionsSchema
     .and(DynatempSchema)
     .and(MirostatSchema);
 
-export type BaseSamplerRequest = z.infer<typeof BaseSamplerRequest>;
+// Define the type from the schema
+export type BaseSamplerRequest = z.infer<typeof BaseSamplerRequestSchema>;
+
+// Apply transforms and expose the type
+export const BaseSamplerRequest = BaseSamplerRequestSchema
+    .transform((obj) => {
+        return forcedSamplerOverrides(obj);
+    });