Change configuration in LLMModule so it can be set before loading model

[msluszniak]

Description

Currently, there is no other way to set configuration in useLLM other than load model first, and then call configure method. This PR make it possible to configure parameters before loading the actual model.

Introduces a breaking change?

• Yes
• No

Type of change

• Bug fix (change which fixes an issue)
• New feature (change which adds functionality)
• Documentation update (improves or adds clarity to existing documentation)
• Other (chores, tests, code style improvements etc.)

Tested on

• iOS
• Android

Testing instructions

Try to run configure on hook returned by useLLM and check that everything works.

For simplicity, I present the example way how to test it inside our library:

• Create the following file in apps/llm/app/my_test/index.tsx:

import { useIsFocused } from '@react-navigation/native';
import React, { useEffect, useState, useRef } from 'react';
import {
  View,
  Text,
  TextInput,
  TouchableOpacity,
  FlatList,
  StyleSheet,
  ActivityIndicator,
  KeyboardAvoidingView,
  Platform,
  SafeAreaView,
} from 'react-native';
import { LLMModule, LLAMA3_2_1B_QLORA } from 'react-native-executorch';

// Define message type for UI
type Message = {
  role: 'user' | 'assistant' | 'system';
  content: string;
};

export default function VoiceChatScreenWrapper() {
  const isFocused = useIsFocused();

  return isFocused ? <LlamaChat /> : null;
}


const LlamaChat = () => {
  const [messages, setMessages] = useState<Message[]>([]);
  const [input, setInput] = useState('');
  const [isModelReady, setIsModelReady] = useState(false);
  const [loadingProgress, setLoadingProgress] = useState(0);
  const [isGenerating, setIsGenerating] = useState(false);

  // Use a ref to keep the LLM instance stable across renders
  const llmRef = useRef<LLMModule | null>(null);

  useEffect(() => {
    // 1. Initialize the LLM Module
    llmRef.current = new LLMModule({
      // Update state whenever history changes (covers both user and bot messages)
      messageHistoryCallback: (updatedMessages) => {
        // We cast this to our Message type (assuming the library returns compatible format)
        setMessages(updatedMessages as Message[]);
      },
      // Optional: Use tokenCallback if you want to trigger haptics or very fine-grained updates
      tokenCallback: (token) => {
        // console.log('New token:', token);
      },
    });

    // 2. Load the model
    const loadModel = async () => {
      try {
        await llmRef.current?.load(LLAMA3_2_1B_QLORA, (progress) => {
          setLoadingProgress(progress);
        });
        setIsModelReady(true);
      } catch (error) {
        console.error('Failed to load model:', error);
      }
    };

    loadModel();

    llmRef.current?.configure({chatConfig: {systemPrompt: "You are extremely enthusiastic chat assistant that is ecstatic about chatting with me."}});

    // 3. Cleanup: Delete model from memory when component unmounts
    return () => {
      console.log('Cleaning up LLM...');
      llmRef.current?.delete();
    };
  }, []);

  const handleSend = async () => {
    if (!input.trim() || !isModelReady || isGenerating) return;

    const userText = input;
    setInput(''); // Clear input immediately
    setIsGenerating(true);

    try {
      // sendMessage automatically updates the history via the callback defined in useEffect
      await llmRef.current?.sendMessage(userText);
    } catch (error) {
      console.error('Error generating response:', error);
    } finally {
      setIsGenerating(false);
    }
  };

  const handleStop = () => {
    llmRef.current?.interrupt();
    setIsGenerating(false);
  };

  // --- Render Helpers ---

  if (!isModelReady) {
    return (
      <View style={styles.centerContainer}>
        <ActivityIndicator size="large" color="#007AFF" />
        <Text style={styles.loadingText}>
          Loading Model... {(loadingProgress * 100).toFixed(0)}%
        </Text>
      </View>
    );
  }

  return (
    <SafeAreaView style={styles.container}>
      <KeyboardAvoidingView
        behavior={Platform.OS === 'ios' ? 'padding' : undefined}
        style={styles.keyboardContainer}
      >
        <FlatList
          data={messages}
          keyExtractor={(_, index) => index.toString()}
          contentContainerStyle={styles.listContent}
          renderItem={({ item }) => (
            <View
              style={[
                styles.bubble,
                item.role === 'user' ? styles.userBubble : styles.botBubble,
              ]}
            >
              <Text style={item.role === 'user' ? styles.userText : styles.botText}>
                {item.content}
              </Text>
            </View>
          )}
        />

        <View style={styles.inputContainer}>
          <TextInput
            style={styles.input}
            placeholder="Ask Llama..."
            value={input}
            onChangeText={setInput}
            editable={!isGenerating}
          />
          
          {isGenerating ? (
            <TouchableOpacity onPress={handleStop} style={styles.stopButton}>
              <Text style={styles.buttonText}>Stop</Text>
            </TouchableOpacity>
          ) : (
            <TouchableOpacity onPress={handleSend} style={styles.sendButton}>
              <Text style={styles.buttonText}>Send</Text>
            </TouchableOpacity>
          )}
        </View>
      </KeyboardAvoidingView>
    </SafeAreaView>
  );
};

const styles = StyleSheet.create({
  container: { flex: 1, backgroundColor: '#F5F5F5' },
  centerContainer: { flex: 1, justifyContent: 'center', alignItems: 'center' },
  loadingText: { marginTop: 10, fontSize: 16, color: '#333' },
  keyboardContainer: { flex: 1 },
  listContent: { padding: 16 },
  bubble: {
    maxWidth: '80%',
    padding: 12,
    borderRadius: 16,
    marginBottom: 10,
  },
  userBubble: {
    alignSelf: 'flex-end',
    backgroundColor: '#007AFF',
    borderBottomRightRadius: 2,
  },
  botBubble: {
    alignSelf: 'flex-start',
    backgroundColor: '#E5E5EA',
    borderBottomLeftRadius: 2,
  },
  userText: { color: '#FFF', fontSize: 16 },
  botText: { color: '#000', fontSize: 16 },
  inputContainer: {
    flexDirection: 'row',
    padding: 10,
    borderTopWidth: 1,
    borderColor: '#DDD',
    backgroundColor: '#FFF',
  },
  input: {
    flex: 1,
    backgroundColor: '#F0F0F0',
    borderRadius: 20,
    paddingHorizontal: 16,
    paddingVertical: 10,
    fontSize: 16,
    marginRight: 10,
  },
  sendButton: {
    backgroundColor: '#007AFF',
    justifyContent: 'center',
    alignItems: 'center',
    paddingHorizontal: 20,
    borderRadius: 20,
  },
  stopButton: {
    backgroundColor: '#FF3B30',
    justifyContent: 'center',
    alignItems: 'center',
    paddingHorizontal: 20,
    borderRadius: 20,
  },
  buttonText: { color: '#FFF', fontWeight: '600' },
});

• Add the following in apps/llm/app/_layout.txs:

+        <Drawer.Screen
+          name="my_test/index"
+          options={{
+            drawerLabel: 'Llama Chat',
+            title: 'Llama Chat',
+            headerTitleStyle: { color: ColorPalette.primary },
+          }}
+        />

• Add the following in apps/llm/app/index.tsx:

+
+          <TouchableOpacity
+          style={styles.button}
+          onPress={() => router.navigate('my_test/')}
+        >
+          <Text style={styles.buttonText}>LLama chat</Text>
+        </TouchableOpacity>

Run llm app and ask about anything. Generation config should work correctly, and now responses of the LLM should be super ecstatic. Now, move this part:

    llmRef.current?.configure({chatConfig: {systemPrompt: "You are extremely enthusiastic chat assistant that is ecstatic about chatting with me."}});

before loading the model and check if everything works correct.

Screenshots

Related issues

Checklist

• I have performed a self-review of my code
• I have commented my code, particularly in hard-to-understand areas
• I have updated the documentation accordingly
• My changes generate no new warnings

Additional notes