Add custom jinja template for text generation

Author: RHeckerIntel

lib/langchain/chain_builder.dart

@@ -17,44 +17,34 @@ String combineDocuments(
     documents.map((final d) => d.pageContent).join(separator);
 
 
-RAGChain buildRAGChain(LLMInference llmInference, Embeddings embeddings, OpenVINOLLMOptions options, List<VectorStore> stores) {
+RAGChain buildRAGChain(LLMInference llmInference, Embeddings embeddings, OpenVINOLLMOptions options, List<VectorStore> stores, BaseChatMemory memory) {
   final retrievers = combineStores(stores);
 
+  final tokenizerConfig = jsonDecode(llmInference.getTokenizerConfig()) as Map<String, dynamic>;
+
   final retrievedDocs = Runnable.fromMap({
     'docs': Runnable.getItemFromMap('question') | retrievers,
     'question': Runnable.getItemFromMap('question'),
   });
 
-  if (stores.isEmpty) {
-    final model = OpenVINOLLM(llmInference, defaultOptions: options.copyWith(applyTemplate: true));
-    final answer = PromptTemplate.fromTemplate('{question}') | model;
-    return RAGChain(retrievedDocs, answer);
-  }
-
-
-  final tokenizerConfig = jsonDecode(llmInference.getTokenizerConfig()) as Map<String, dynamic>;
-
-  final hasChatTemplate = tokenizerConfig.containsKey("chat_template");
-
-  // if chat template, otherwise
-  final promptTemplate = hasChatTemplate
-    ? JinjaPromptTemplate.fromTemplateConfig(tokenizerConfig)
-    : ChatPromptTemplate.fromTemplate('''
-Answer the question based only on the following context without specifically naming that it's from that context:
-{context}
-
-Question: {question}
-''');
+  final promptTemplate = JinjaPromptTemplate.fromTemplateConfig(tokenizerConfig);
 
   final finalInputs = Runnable.fromMap({
     'context': Runnable.getItemFromMap<List<Document>>('docs') |
         Runnable.mapInput<List<Document>, String>(combineDocuments),
     'question': Runnable.getItemFromMap('question'),
+    'history': Runnable.getItemFromMap('question') | Runnable.mapInput((_) async {
+      final m = await memory.loadMemoryVariables();
+      return m['history'];
+    }),
   });
   final model = OpenVINOLLM(llmInference, defaultOptions: options.copyWith(applyTemplate: false));
 
   final answer = finalInputs | promptTemplate | model;
 
+  finalInputs.invoke({'docs': List<Document>.from([]), 'question': "What is the color of the sun?"}).then(print);
+
+
   return RAGChain(retrievedDocs, answer);
 }
 

lib/langchain/jinja_prompt_template.dart

@@ -5,6 +5,14 @@
 import 'package:jinja/jinja.dart';
 import 'package:langchain/langchain.dart';
 
+const textGenerationTemplate = """
+{% for message in messages %}{%- if message['role'] == 'system' %}{{message['content']}}
+
+Question:
+{%- endif %}{% endfor %}
+{% for message in messages %}{%- if message['role'] == 'user' %}{{message['content']}}{%- endif %}{% endfor %}
+""";
+
 final class JinjaPromptTemplate extends BaseChatPromptTemplate {
   final Template jinjaTemplate;
 
@@ -19,7 +27,9 @@ final class JinjaPromptTemplate extends BaseChatPromptTemplate {
   });
 
   factory JinjaPromptTemplate.fromTemplateConfig(Map<String, dynamic> chatTemplateConfig, [Set<String> inputVariables = const {}]) {
-    final chatTemplate = chatTemplateConfig["chat_template"];
+    final chatTemplate = chatTemplateConfig.containsKey("chat_template")
+      ? chatTemplateConfig["chat_template"]
+      : textGenerationTemplate;
     final env = Environment();
     final template = env.fromString(chatTemplate);
 
@@ -46,10 +56,23 @@ final class JinjaPromptTemplate extends BaseChatPromptTemplate {
 
   @override
   PromptValue formatPrompt(final InputValues values) {
-    final messages =[
-      {"role": "system", "content": "Answer the question based on some info:\n ${values['context']}"},
-      {"role": "user", "content": values['question']},
-    ];
+    List<Map<String, dynamic>> messages = [];
+    if (values.containsKey('history')) {
+      for (final message in values['history']) {
+        if (message is AIChatMessage) {
+          messages.add({"role": "assistant", "content": message.contentAsString});
+        }
+        if (message is HumanChatMessage) {
+          messages.add({"role": "user", "content": message.contentAsString});
+        }
+      }
+    }
+    if (values.containsKey('context') && values['context'] != "") {
+      messages.add({"role": "system", "content": "Answer the question based on some info:\n ${values['context']}"});
+    }
+    if (values.containsKey('question')) {
+      messages.add({"role": "user", "content": values['question']});
+    }
 
     return PromptValue.string(jinjaTemplate.render(
       {

lib/providers/text_inference_provider.dart

@@ -62,6 +62,8 @@ class TextInferenceProvider extends ChangeNotifier {
   String? get device => _device;
   Metrics? get metrics => _messages.lastOrNull?.metrics;
 
+  final memory = ConversationBufferMemory(returnMessages: true);
+
   final List<UserFile> _userFiles = [];
 
   Future<void> addUserFiles(List<UserFile> files ) async {
@@ -199,10 +201,11 @@ class TextInferenceProvider extends ChangeNotifier {
       stores.add(ObjectBoxStore(embeddings: embeddingsModel!, group: knowledgeGroup!));
     }
 
-    final chain = buildRAGChain(_inference!, embeddingsModel!, OpenVINOLLMOptions(temperature: temperature, topP: topP), stores);
+    final chain = buildRAGChain(_inference!, embeddingsModel!, OpenVINOLLMOptions(temperature: temperature, topP: topP), stores, memory);
     final input = await chain.documentChain.invoke({"question": message}) as Map;
-    print(input);
-    final docs = List<String>.from(input["docs"].map((Document doc) => doc.metadata["source"]).toSet());
+    final docs = input.containsKey("docs")
+      ? List<String>.from(input["docs"].map((Document doc) => doc.metadata["source"]).toSet())
+      : null;
 
     String modelOutput = "";
     Metrics? metrics;
@@ -216,6 +219,11 @@ class TextInferenceProvider extends ChangeNotifier {
       onToken(token);
     }
 
+    memory.saveContext(
+      inputValues: {'input': message},
+      outputValues: {'output': modelOutput},
+    );
+
     if (_messages.isNotEmpty) {
       _messages.add(Message(Speaker.assistant, modelOutput, metrics, DateTime.now(), sources: docs));
     }
@@ -250,6 +258,7 @@ class TextInferenceProvider extends ChangeNotifier {
   void reset() {
     _inference?.forceStop();
     _inference?.clearHistory();
+    memory.clear();
     for (final file in _userFiles) {
       final ids = file.documents.map((p) => p.id).whereType<String>().toList();
       store?.delete(ids: ids);

openvino_bindings/src/llm/llm_inference.cc

@@ -30,9 +30,10 @@ ov::genai::DecodedResults LLMInference::prompt(std::string message, bool apply_t
     history.push_back({{"role", "user"}, {"content", message}});
     _stop = false;
 
-    auto prompt = (apply_template && has_chat_template()
-        ? pipe.get_tokenizer().apply_chat_template(history, true)
-        : message);
+    //auto prompt = (apply_template && has_chat_template()
+    //    ? pipe.get_tokenizer().apply_chat_template(history, true)
+    //    : message);
+    auto prompt = message;
 
     ov::genai::GenerationConfig config;
     config.max_new_tokens = 1000;