-
-
-
-
-
-
- Normal Gemma
+
+ Strength
+
+
+
{
+ setStrengthLevel(value[0]);
+ }}
+ className="relative mt-1 flex h-5 w-full flex-1 cursor-pointer items-center"
+ >
+
+
-
- {!isTuning && steeredChatMessages.length === 0 && (
-
- Hey, {`I'm normal Gemma!`}
-
- Get started below.
-
- )}
-
- {isTuning &&
}
-
-
-
-
-
- {!isTuning && steeredChatMessages.length === 0 && (
-
- Hey, {`I'm steered Gemma!`}
-
- Get started below.
-
- )}
-
- {isTuning &&
}
-
-
-
-
-
-
-
-
{
- if (e.key === 'Enter' && !e.shiftKey && !isTuning) {
- sendChat();
- e.preventDefault();
- }
- }}
- onChange={(e) => {
- // if it's return, submit
- if (e.target.value.indexOf('\n') === -1) {
- setTypedInText(e.target.value);
- }
- }}
- required
- placeholder="Ask Gemma something..."
- className="mt-0 w-full flex-1 resize-none rounded-full border border-slate-300 px-5 py-3.5 pr-10 text-left text-xs font-medium text-slate-800 placeholder-slate-400 shadow transition-all focus:border-slate-300 focus:shadow focus:outline-none focus:ring-0 disabled:bg-slate-200 sm:text-[13px]"
- />
-
-
{
- if (!isTuning) {
- sendChat();
- }
- }}
- className={`h-8 w-8 rounded-full ${isTuning ? 'bg-slate-400' : 'bg-gBlue hover:bg-gBlue/80'
- } p-1.5 text-white`}
- />
+
+
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
- Feature
-
-
- 0 ? featurePresets[0].name : 'Loading...'}
- value={
- // find where selectedFeature matches a feature in the presets
- featurePresets.find((p) =>
- selectedFeatures.find((f) =>
- p.features.find((fp) => fp.modelId === f.modelId && fp.layer === f.layer && fp.index === f.index),
- ),
- )?.name || 'Loading...'
- }
- onValueChange={(presetName) => {
- const feat = featurePresets.find((p) => p.name === presetName);
- setSelectedFeatures(feat?.features || []);
- }}
- >
-
-
-
-
-
-
-
-
-
-
-
-
- {featurePresets.length > 0 ? (
- featurePresets.map((preset) => (
-
-
- {preset.name}
-
-
- ))
- ) : (
-
-
- Loading...
-
-
- )}
-
-
-
-
-
-
-
-
-
-
- Strength
-
-
-
{
- setStrengthLevel(value[0]);
- }}
- className="relative mt-1 flex h-5 w-full flex-1 cursor-pointer items-center"
- >
-
-
-
-
-
- 0️⃣
}>0️⃣ No Steering
-
-
{' '}
-
-
- ⚖️
}>⚖️ Medium Strength
-
-
{' '}
-
-
- 🤯
}>🤯 Quite A Bit
-
-
{' '}
-
-
-
-
- {`+${strengthLevel.toFixed(1)}x`}
-
-
- {/*
*/}
-
-
- {showMoreOptions && (
-
-
- Steering Method
-
- {/*
{
- alert('Activation not available yet.');
- return;
- setSteeringMethod(value as SteeringMethod);
- }}
- aria-label="steering method"
- >
-
- Activation
-
-
- SAE
-
- */}
-
-
- Advanced
-
-
-
-
- Strength Multiple
+
+
-
{
- if (
- parseFloat(e.target.value) < 0 ||
- parseFloat(e.target.value) > STEER_STRENGTH_MULTIPLIER_MAX
- ) {
- alert(`Strength multiplier must be >= 0 and <= ${STEER_STRENGTH_MULTIPLIER_MAX}`);
- } else {
- setStrMultiple(parseFloat(e.target.value));
- }
- }}
- className="max-w-[80px] flex-1 rounded-md border-green-400 py-1 text-center text-xs text-green-800"
- value={strMultiple}
- />
-
-
- Steer Special Tokens
+
-
- {
- setSteerSpecialTokens(e.target.checked);
- }}
- type="checkbox"
- checked={steerSpecialTokens}
- className="h-5 w-5 cursor-pointer rounded border-green-400 bg-green-100 py-1 text-center text-xs text-green-800 checked:bg-green-600 checked:text-white"
- />
+
{' '}
+
+
+ ⚖️
}>⚖️ Medium Strength
-
-
-
- )}
-
- {showMoreOptions && (
-
-
-
- Generation Settings
-
-
-
-
- Num Tokens
+
{' '}
+
-
{
- if (parseInt(e.target.value, 10) > STEER_N_COMPLETION_TOKENS_MAX) {
- alert(
- `Due to compute constraints, the current allowed max tokens is: ${STEER_N_COMPLETION_TOKENS_MAX}`,
- );
- } else {
- setSteerTokens(parseInt(e.target.value, 10));
- }
- }}
- className="max-w-[80px] flex-1 rounded-md border-amber-400 py-1 text-center text-xs text-amber-800"
- value={steerTokens}
- />
-
-
-
- Temperature
+
+ 🤯
}>🤯 Quite A Bit
-
{
- if (parseFloat(e.target.value) > STEER_TEMPERATURE_MAX || parseFloat(e.target.value) < 0) {
- alert(`Temperature must be >= 0 and <= ${STEER_TEMPERATURE_MAX}`);
- } else {
- setTemperature(parseFloat(e.target.value));
- }
- }}
- className="max-w-[80px] flex-1 rounded-md border-amber-400 py-1 text-center text-xs text-amber-800"
- value={temperature}
- />
-
-
-
- Freq Penalty
-
-
{
- if (parseFloat(e.target.value) > 2 || parseFloat(e.target.value) < -2) {
- alert('Freq penalty must be >= -2 and <= 2');
- } else {
- setFreqPenalty(parseFloat(e.target.value));
- }
- }}
- className="max-w-[80px] flex-1 rounded-md border-amber-400 py-1 text-center text-xs text-amber-800"
- value={freqPenalty}
- />
-
-
-
Seed
-
{
- if (parseInt(e.target.value, 10) > 100000000 || parseInt(e.target.value, 10) < -100000000) {
- alert('Seed must be >= -100000000 and <= 100000000');
- } else {
- setSeed(parseInt(e.target.value, 10));
- }
- }}
- className="max-w-[80px] flex-1 rounded-md border-amber-400 py-1 text-center text-xs text-amber-800 disabled:bg-amber-200 disabled:text-amber-400"
- value={seed}
- />
-
-
-
-
- )}
- {showOptionsButton ? (
- <>
-
- {/*
-
*/}
-
- ) : (
-
- )}
+
+
+
+ {`+${strengthLevel.toFixed(1)}x`}
+
+
+
-
+
+ );
+
+ // const handleToggleMoreOptions = useCallback(() => { // Now from hook
+ // setShowMoreOptions((prev) => !prev);
+ // }, [setShowMoreOptions]);
+
+ return (
+
+ }
+ showOptionsButton={showOptionsButton} // Prop
+ >
+ {steeringControlsJsx}
+
);
}
diff --git a/apps/webapp/components/tools/feature-discovery-assistant.tsx b/apps/webapp/components/tools/feature-discovery-assistant.tsx
new file mode 100644
index 000000000..a33194cc9
--- /dev/null
+++ b/apps/webapp/components/tools/feature-discovery-assistant.tsx
@@ -0,0 +1,417 @@
+/* eslint-disable jsx-a11y/label-has-associated-control -- customComponent doc examples don't work :( */
+
+'use client';
+
+import ModelSelector from '@/components/feature-selector/model-selector';
+import { useGlobalContext } from '@/components/provider/global-provider';
+import { Button } from '@/components/shadcn/button';
+import { SearchExplanationsResponse } from '@/lib/utils/general';
+import {
+ SteerFeature,
+ STEER_TEMPERATURE,
+ STEER_FREQUENCY_PENALTY,
+ STEER_SEED,
+ STEER_STRENGTH_MULTIPLIER,
+ STEER_SPECIAL_TOKENS,
+ ChatMessage as SteerChatMessage, // Renaming to avoid conflict with local ChatMessage
+} from '@/lib/utils/steer';
+import { callSteerChatApi } from '@/lib/utils/steer-api'; // Import the new utility
+// import { ExplanationWithPartialRelations } from '@/prisma/generated/zod'; // No longer directly used in this component's props/state
+import { useState, useEffect, useRef, useCallback } from 'react';
+
+// Define a type for our structured feature suggestion
+type SuggestedSteerFeature = SteerFeature & {
+ originalExplanation: string; // To show the source explanation
+ score?: number; // Placeholder for potential future sorting/relevance score
+ isSelected: boolean; // To manage selection state
+};
+
+// Define FeatureSteeringResult type
+// This should extend SuggestedSteerFeature to include steering-specific results
+type FeatureSteeringResult = SuggestedSteerFeature & {
+ steeredText?: string;
+ steeringError?: string;
+ isSteeringLoading?: boolean;
+};
+
+export default function FeatureDiscoveryAssistant() {
+ const { getDefaultModel, getInferenceEnabledModels, showToastServerError } = useGlobalContext();
+ const [modelId, setModelId] = useState
(getDefaultModel()?.id || '');
+ const [testPrompt, setTestPrompt] = useState('');
+ const [featureQuery, setFeatureQuery] = useState('');
+ const [isLoading, setIsLoading] = useState(false); // For the main search
+ const [suggestedFeatures, setSuggestedFeatures] = useState([]);
+ // steeringResults will store the outcome for each of the 5 suggestedFeatures
+ const [steeringResults, setSteeringResults] = useState([]);
+ const [hasSteered, setHasSteered] = useState(false); // Tracks if "Steer Selected Features" has been clicked
+
+ // State for auto-steer timer
+ const [autoSteerTimerId, setAutoSteerTimerId] = useState(null);
+ const activeTimerIdRef = useRef(null); // To robustly check active timer in setTimeout
+ const [userInteracted, setUserInteracted] = useState(false);
+ const [triggerAutoSteer, setTriggerAutoSteer] = useState(false);
+
+
+ const handleSearch = async () => {
+ // Reset timer and interaction states on new search
+ if (autoSteerTimerId) {
+ clearTimeout(autoSteerTimerId);
+ }
+ setAutoSteerTimerId(null);
+ if (activeTimerIdRef.current) {
+ clearTimeout(activeTimerIdRef.current);
+ activeTimerIdRef.current = null;
+ }
+ setUserInteracted(false);
+ setHasSteered(false); // Also reset hasSteered as per instructions
+
+ if (!modelId || !testPrompt.trim() || !featureQuery.trim()) {
+ console.warn('Please select a model, enter a test prompt, and enter a feature name to search.');
+ return;
+ }
+ setIsLoading(true);
+ setSuggestedFeatures([]);
+ setSteeringResults([]);
+ // setHasSteered(false); // Already reset above
+
+ try {
+ const response = await fetch('/api/explanation/search-model', {
+ method: 'POST',
+ headers: { 'Content-Type': 'application/json' },
+ body: JSON.stringify({ query: featureQuery, modelId, offset: 0 }),
+ });
+
+ if (!response.ok) throw new Error(`API Error: ${response.status} ${response.statusText}`);
+ const data = (await response.json()) as SearchExplanationsResponse; // Type assertion for data
+
+ if (data.results && data.results.length > 0) {
+ const topResults = data.results.slice(0, 5);
+ const formattedFeatures: SuggestedSteerFeature[] = topResults.map((exp, index) => {
+ const { neuron } = exp; // exp is an ExplanationWithPartialRelations here
+ if (!neuron || !neuron.index) {
+ console.warn('Explanation result missing neuron data:', exp); return null;
+ }
+ return {
+ modelId: neuron.modelId,
+ layer: neuron.layer,
+ index: parseInt(neuron.index, 10),
+ explanation: `Feature related to: "${exp.description}" (Original model: ${neuron.modelId}, Layer: ${neuron.layer}, Index: ${neuron.index})`,
+ strength: neuron.maxActApprox || 1,
+ originalExplanation: exp.description,
+ isSelected: index < 2, // Default select top 2
+ };
+ }).filter(Boolean) as SuggestedSteerFeature[];
+ setSuggestedFeatures(formattedFeatures);
+
+ // Initialize steeringResults for all 5 discovered features
+ const initialSteeringData: FeatureSteeringResult[] = formattedFeatures.map(feature => ({
+ ...feature,
+ isSteeringLoading: false, // Not loading initially
+ steeredText: undefined,
+ steeringError: undefined,
+ }));
+ setSteeringResults(initialSteeringData);
+
+ // Start auto-steer timer if features were found
+ if (formattedFeatures.length > 0) {
+ const newTimerId = setTimeout(() => {
+ if (!userInteracted && activeTimerIdRef.current === newTimerId) {
+ console.log('Auto-steering: Timer elapsed without user interaction.');
+ // Auto-select top 3 features
+ setSuggestedFeatures(prevFeatures =>
+ prevFeatures.map((feature, index) => ({
+ ...feature,
+ isSelected: index < 3,
+ }))
+ );
+ setSteeringResults(prevResults =>
+ prevResults.map((result, index) => ({
+ ...result,
+ isSelected: index < 3,
+ }))
+ );
+ setTriggerAutoSteer(true); // Trigger effect to call handleSteerSelectedFeatures
+ }
+ }, 5000);
+ activeTimerIdRef.current = newTimerId;
+ setAutoSteerTimerId(newTimerId);
+ }
+ } else {
+ alert('No features found for your query.');
+ }
+ } catch (error) {
+ console.error('Failed to search features:', error);
+ showToastServerError();
+ alert(`Failed to search features. ${error instanceof Error ? error.message : 'Please check console for details.'}`);
+ } finally {
+ setIsLoading(false);
+ }
+ };
+
+ const handleSteerSelectedFeatures = useCallback(async () => {
+ const selectedForSteering = suggestedFeatures.filter(f => f.isSelected);
+ if (selectedForSteering.length === 0) {
+ alert('Please select at least one feature to steer.');
+ return;
+ }
+ setHasSteered(true);
+
+ // Set loading state only for selected features in the main steeringResults array
+ setSteeringResults(prevGlobalResults =>
+ prevGlobalResults.map(res => {
+ const isSelectedForThisOperation = selectedForSteering.some(
+ sf => sf.modelId === res.modelId && sf.layer === res.layer && sf.index === res.index
+ );
+ return isSelectedForThisOperation
+ ? { ...res, isSteeringLoading: true, steeredText: undefined, steeringError: undefined }
+ : res;
+ })
+ );
+
+ const steeringPromises = selectedForSteering.map(async (featureToSteer): Promise => {
+ try {
+ const apiFeature = {
+ modelId: featureToSteer.modelId,
+ layer: featureToSteer.layer,
+ index: featureToSteer.index,
+ strength: (featureToSteer.strength || 1) * 1.5,
+ explanation: featureToSteer.explanation,
+ };
+ const modelToSteer = (modelId !== 'gemma-2-2b') ? modelId : 'gemma-2-2b-it';
+ const steeringData = await callSteerChatApi({
+ modelId: modelToSteer,
+ defaultChatMessages: [{ role: 'user', content: testPrompt }] as SteerChatMessage[],
+ steeredChatMessages: [{ role: 'user', content: testPrompt }] as SteerChatMessage[],
+ features: [apiFeature],
+ temperature: STEER_TEMPERATURE, nTokens: 96, freqPenalty: STEER_FREQUENCY_PENALTY,
+ seed: STEER_SEED, strengthMultiplier: STEER_STRENGTH_MULTIPLIER, steerSpecialTokens: STEER_SPECIAL_TOKENS,
+ });
+ const assistantMessage = steeringData?.STEERED?.chatTemplate?.findLast(msg => msg.role === 'model' || msg.role === 'assistant')?.content;
+ return { ...featureToSteer, steeredText: assistantMessage || 'No response text.', steeringError: undefined, isSteeringLoading: false };
+ } catch (error) {
+ const errorMessage = error instanceof Error ? error.message : 'Steering request failed';
+ return { ...featureToSteer, steeredText: undefined, steeringError: errorMessage, isSteeringLoading: false };
+ }
+ });
+
+ const settledResults = await Promise.allSettled(steeringPromises);
+
+ const updatedResultsFromPromises: FeatureSteeringResult[] = settledResults.map((settledResult, index) => {
+ if (settledResult.status === 'fulfilled') {
+ return settledResult.value;
+ }
+ const { reason } = settledResult;
+ console.error('Unexpected promise rejection post-settled:', reason);
+ const errorMessage = reason instanceof Error ? reason.message : String(reason);
+ const originalFeatureForRejection = selectedForSteering[index];
+ return {
+ ...originalFeatureForRejection,
+ steeredText: undefined,
+ steeringError: `Unexpected error: ${errorMessage}`,
+ isSteeringLoading: false,
+ };
+ });
+
+ setSteeringResults(prevGlobalResults =>
+ prevGlobalResults.map(existingGlobalFeatureResult => {
+ const newlySteeredData = updatedResultsFromPromises.find(
+ nr => nr.modelId === existingGlobalFeatureResult.modelId &&
+ nr.layer === existingGlobalFeatureResult.layer &&
+ nr.index === existingGlobalFeatureResult.index
+ );
+ if (newlySteeredData) {
+ return {
+ ...existingGlobalFeatureResult,
+ ...newlySteeredData,
+ isSteeringLoading: false,
+ };
+ }
+ return existingGlobalFeatureResult;
+ })
+ );
+ }, [suggestedFeatures, modelId, testPrompt, showToastServerError]);
+
+ useEffect(() => {
+ if (triggerAutoSteer) {
+ handleSteerSelectedFeatures();
+ setTriggerAutoSteer(false); // Reset trigger
+ }
+ }, [triggerAutoSteer, handleSteerSelectedFeatures]);
+
+ // Cleanup timer on component unmount
+ useEffect(() => () => {
+ if (activeTimerIdRef.current) {
+ clearTimeout(activeTimerIdRef.current);
+ }
+ // Also clear state timer if it matches, though ref is primary for active timer
+ if (autoSteerTimerId) {
+ clearTimeout(autoSteerTimerId);
+ }
+ }, [autoSteerTimerId]); // Include autoSteerTimerId to re-run cleanup if it changes, though ref is main
+
+ const handleCheckboxChange = (changedIndex: number) => {
+ if (!userInteracted) {
+ setUserInteracted(true);
+ }
+ if (autoSteerTimerId) {
+ console.log('Auto-steering: User interacted, clearing timer.');
+ clearTimeout(autoSteerTimerId);
+ setAutoSteerTimerId(null);
+ }
+ if (activeTimerIdRef.current) {
+ clearTimeout(activeTimerIdRef.current);
+ activeTimerIdRef.current = null;
+ }
+
+ setSuggestedFeatures(prevFeatures =>
+ prevFeatures.map((feature, idx) =>
+ idx === changedIndex ? { ...feature, isSelected: !feature.isSelected } : feature
+ )
+ );
+ };
+
+ return (
+
+
Feature Discovery Assistant
+
+ {/* Inputs Section */}
+
+
+
+
+
+
+
+
+
+
E.g., "social isolation", "positive sentiment", "code generation"
+
+
+
+
+
+ {/* Discovered Features Section */}
+ {suggestedFeatures.length > 0 && (
+
+
Discovered Features:
+
+ {suggestedFeatures.map((feature, idx) => {
+ const steeringResultForCard = steeringResults.find(
+ sr => sr.modelId === feature.modelId && sr.layer === feature.layer && sr.index === feature.index
+ );
+
+ const cardClasses = [
+ "w-full sm:w-1/2 lg:w-1/3 xl:w-1/4 p-2 mb-4 rounded-md bg-slate-50 shadow-sm transition-all duration-300 ease-in-out",
+ ];
+ if (hasSteered && !feature.isSelected) {
+ cardClasses.push("filter grayscale opacity-50");
+ }
+
+ return (
+
+
+
Feature #{idx + 1}
+ handleCheckboxChange(idx)}
+ className="form-checkbox h-5 w-5 text-emerald-600 transition duration-150 ease-in-out"
+ />
+
+
+ Derived from: "{feature.originalExplanation}"
+
+
+ modelId: '{feature.modelId}', layer: '{feature.layer}', index: {feature.index}, strength: {feature.strength}
+
+
+ {/* Steering Result Display within Card */}
+ {steeringResultForCard && (steeringResultForCard.isSteeringLoading || (hasSteered && feature.isSelected)) && (
+
+ {steeringResultForCard.isSteeringLoading ? (
+
Steering in progress...
+ ) : steeringResultForCard.steeringError ? (
+
+
Steering Error:
+
{steeringResultForCard.steeringError}
+
+ ) : steeringResultForCard.steeredText ? (
+
+
Steered Output:
+
{steeringResultForCard.steeredText}
+
+ ) : (hasSteered && feature.isSelected) ? (
+
Steering completed, but no text was returned.
+ ) : null}
+
+ )}
+
+ );
+ })}
+
+
+ Select features and click "Steer Selected Features" to test their effect.
+
+
+ )}
+
+ {/* "Steer Selected Features" Button */}
+ {suggestedFeatures.length > 0 && (
+
+ )}
+