Merge pull request #95 from ihatenumbers/banned_tokens

Add Banned Tokens support

mikupad.html

@@ -4639,6 +4639,7 @@
 	dryAllowedLength: 2,
 	dryPenaltyRange: 1024,
 	drySequenceBreakers: "[\"\\n\", \":\", \"\\\"\", \"*\"]",
+	bannedTokens: "[]",
 	stoppingStrings: "[]",
 	ignoreEos: false,
 	openaiPresets: false,
@@ -4908,6 +4909,8 @@
 	const [dryPenaltyRange, setDryPenaltyRange] = useSessionState('dryPenaltyRange', defaultPresets.dryPenaltyRange);
 	const [drySequenceBreakers, setDrySequenceBreakers] = useSessionState('drySequenceBreakers', defaultPresets.drySequenceBreakers);
 	const [drySequenceBreakersError, setDrySequenceBreakersError] = useState(undefined);
+	const [bannedTokens, setBannedTokens] = useSessionState('bannedTokens', defaultPresets.bannedTokens);
+	const [bannedTokensError, setBannedTokensError] = useState(undefined);
 	const [ignoreEos, setIgnoreEos] = useSessionState('ignoreEos', defaultPresets.ignoreEos);
 	const [openaiPresets, setOpenaiPresets] = useSessionState('openaiPresets', defaultPresets.openaiPresets);
 	const [rejectedAPIKey, setRejectedAPIKey] = useState(false);
@@ -5484,6 +5487,9 @@
 							dry_penalty_last_n: dryPenaltyRange,
 							dry_sequence_breakers: JSON.parse(drySequenceBreakers),
 						}: {}),
+						...(enabledSamplers.includes('ban_tokens') ? {
+							banned_tokens: JSON.parse(bannedTokens),
+						}: {}),
 					} : {})
 				}),
 				n_predict: maxPredictTokens,
@@ -5653,6 +5659,14 @@
 		}
 	}, [drySequenceBreakers]);
 
+	useEffect(() => {
+		try {JSON.parse(bannedTokens);
+			setBannedTokensError(undefined);
+		} catch (e) {
+			setBannedTokensError(e.toString());
+		}
+	}, [bannedTokens]);
+
 	useEffect(() => {
 		if (showProbsMode === -1 || tokenHighlightMode === -1)
 			return;
@@ -6396,6 +6410,12 @@
 						onValueChange=${(v) => enabledSamplers.indexOf('dry') === -1
 											  ? setEnabledSamplers((es) => [...es, 'dry'])
 											  : setEnabledSamplers((es) => es.filter((s) => s !== 'dry'))}/>
+					<${Checkbox} label="Banned Strings"
+						disabled=${!!cancel}
+						value=${enabledSamplers.includes('ban_tokens')}
+						onValueChange=${(v) => enabledSamplers.indexOf('ban_tokens') === -1
+											  ? setEnabledSamplers((es) => [...es, 'ban_tokens'])
+											  : setEnabledSamplers((es) => es.filter((s) => s !== 'ban_tokens'))}/>
 					<${Checkbox} label="Top K"
 						disabled=${!!cancel}
 						value=${enabledSamplers.includes('top_k')}
@@ -6507,6 +6527,13 @@
 									readOnly=${!!cancel}
 									value=${drySequenceBreakers}
 									onValueChange=${setDrySequenceBreakers}/>`}
+							${enabledSamplers.includes('ban_tokens') && html`
+								<${InputBox} label="Banned Strings (JSON array)" type="text" pattern="^\\[.*?\\]$"
+									className="${bannedTokensError ? 'rejected' : ''}"
+									tooltip="${bannedTokensError ? bannedTokensError : ''}"
+									readOnly=${!!cancel}
+									value=${bannedTokens}
+									onValueChange=${setBannedTokens}/>`}
 						`}
 						${(enabledSamplers.includes('top_k') || enabledSamplers.includes('top_p') || enabledSamplers.includes('min_p')) && html`
 							<div className="hbox">
@@ -6603,7 +6630,7 @@
 				<button
 					title="Run next prediction (Ctrl + Enter)"
 					className=${cancel && !sessionEndpointConnecting ? (predictStartTokens === tokens ? 'processing' : 'completing') : ''}
-					disabled=${!!cancel || stoppingStringsError || drySequenceBreakersError}
+					disabled=${!!cancel || stoppingStringsError || drySequenceBreakersError || bannedTokensError}
 					onClick=${() => predict()}>
 					Predict
 				</button>